Git Worktree Creator with Submodule Support - Automates creation of worktrees for monolith repositories

create_worktree.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.8"
# dependencies = [
#     "rich>=13.7.0",
#     "questionary>=2.0.1",
# ]
# ///

"""
Git Worktree Creator with Submodule Support

This script automates the creation of Git worktrees for monolith repositories
with intelligent handling of submodules. It's designed to work with Python 3.8+
and uses inline script dependencies via PEP 723.

PURPOSE:
--------
When working with a monolith repository containing multiple repositories as submodules,
you often need to create worktrees with different branch configurations for testing 
features that span multiple repositories. This script simplifies that process.

FEATURES:
---------
- Interactive branch selection for main repository and all submodules
- Smart branch filtering to show common branches first (main, master, develop, etc.)
- Search functionality for finding specific branches among hundreds
- Automatic detection of current branches in each submodule
- Proper submodule initialization in the new worktree
- Handles branch conflicts (when a branch is already checked out elsewhere)
- Copies environment files to the new worktree

USAGE:
------
Run directly with Python 3.8+:
    python create_worktree.py

Or using uv (recommended):
    uv run create_worktree.py

WORKFLOW:
---------
1. Script prompts for a worktree name (default: {current-branch}-worktree)
2. Select branch for main repository
3. For each submodule, select desired branch
4. Displays configuration summary for confirmation
5. Creates worktree one directory above current location
6. Initializes all submodules and checks out selected branches
7. Copies environment files if they exist

Author: Diversio Engineering Team
License: MIT
"""

import subprocess
import sys
from pathlib import Path
from typing import Dict, Optional, List
import shutil

try:
    import questionary
    from rich.console import Console
    from rich.table import Table
    from rich.panel import Panel
except ImportError:
    print("Error: Required dependencies not installed.")
    print("Install with: pip install rich questionary")
    sys.exit(1)

console = Console()


def run_command(
    cmd: List[str], cwd: Optional[Path] = None, check: bool = True
) -> subprocess.CompletedProcess:
    """Run a command and return the result."""
    return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, check=check)


def get_current_branch() -> str:
    """Get the current branch name."""
    result = run_command(["git", "symbolic-ref", "--short", "HEAD"], check=False)
    if result.returncode != 0:
        # Try to get commit hash if in detached state
        result = run_command(["git", "rev-parse", "--short", "HEAD"])
        return f"detached-{result.stdout.strip()}"
    return result.stdout.strip()


def get_submodules() -> Dict[str, str]:
    """Get all submodules and their current branches."""
    submodules = {}

    # Get submodule paths
    result = run_command(
        ["git", "config", "--file", ".gitmodules", "--get-regexp", "path"]
    )

    for line in result.stdout.strip().split("\n"):
        if line:
            _, path = line.split(None, 1)
            # Try to get the current branch for each submodule
            try:
                branch_result = run_command(
                    ["git", "symbolic-ref", "--short", "HEAD"],
                    cwd=Path(path),
                    check=False,
                )
                if branch_result.returncode == 0:
                    branch = branch_result.stdout.strip()
                else:
                    # Might be in detached HEAD state
                    branch = "main"  # Default fallback
            except Exception:
                branch = "main"

            submodules[path] = branch

    return submodules


def get_remote_branches(repo_path: Optional[Path] = None) -> List[str]:
    """Get all remote branches for a repository."""
    result = run_command(["git", "branch", "-r"], cwd=repo_path, check=False)
    if result.returncode != 0:
        return ["main", "master", "develop"]  # Fallback options

    branches = []
    for line in result.stdout.strip().split("\n"):
        line = line.strip()
        if line and "->" not in line:  # Skip HEAD pointers
            branch = line.replace("origin/", "")
            if branch not in branches:
                branches.append(branch)

    return branches if branches else ["main", "master", "develop"]


def get_common_branches(all_branches: List[str], current_branch: str) -> List[str]:
    """Get common/important branches prioritized."""
    priority_branches = ["main", "master", "develop", "staging", "production", "release"]
    common_branches = []

    # Add current branch first if it exists
    if current_branch in all_branches:
        common_branches.append(current_branch)

    # Add priority branches that exist
    for branch in priority_branches:
        if branch in all_branches and branch not in common_branches:
            common_branches.append(branch)

    # Add the option to search for more
    if len(all_branches) > len(common_branches):
        common_branches.append("🔍 Search for a different branch...")

    return common_branches


def copy_environment_files(source_dir: Path, target_dir: Path) -> None:
    """Copy environment files from source to target directory."""
    # Define common environment file patterns
    env_patterns = [
        ".env",
        ".env.local",
        ".env.development",
        ".env.production",
        "*.env",
    ]
    
    # Find all env files recursively (up to 2 levels deep)
    env_files = []
    for pattern in env_patterns:
        # Root level
        env_files.extend(source_dir.glob(pattern))
        # One level deep
        env_files.extend(source_dir.glob(f"*/{pattern}"))
        # Two levels deep  
        env_files.extend(source_dir.glob(f"*/*/{pattern}"))
    
    copied_files = []
    for source_file in env_files:
        if source_file.is_file():
            # Calculate relative path
            try:
                relative_path = source_file.relative_to(source_dir)
                target_file = target_dir / relative_path
                
                # Ensure target directory exists
                target_file.parent.mkdir(parents=True, exist_ok=True)
                
                # Copy the file
                shutil.copy2(source_file, target_file)
                copied_files.append(str(relative_path))
                console.print(f"[green]✓[/green] Copied {relative_path}")
            except Exception as e:
                console.print(f"[yellow]Warning: Could not copy {source_file.name}: {e}[/yellow]")
    
    if copied_files:
        console.print(f"\n[bold green]Copied {len(copied_files)} environment file(s)![/bold green]")
        console.print("[dim]These files contain configuration needed to run the applications.[/dim]")
    else:
        console.print(f"\n[dim]No environment files found to copy.[/dim]")


def select_branch_interactive(
    repo_name: str,
    all_branches: List[str],
    current_branch: str
) -> str:
    """Interactive branch selection with search capability."""
    common_branches = get_common_branches(all_branches, current_branch)
    
    selected = questionary.select(
        f"Select branch for {repo_name}:",
        choices=common_branches,
        default=current_branch if current_branch in common_branches else common_branches[0],
    ).ask()

    if not selected:
        console.print("[red]Cancelled.[/red]")
        sys.exit(0)

    # Handle search option
    if selected == "🔍 Search for a different branch...":
        search_query = questionary.text(
            f"Enter branch name or pattern to search in {repo_name}:",
            default=""
        ).ask()

        if not search_query:
            console.print("[red]Cancelled.[/red]")
            sys.exit(0)

        # Filter branches
        filtered_branches = [
            b for b in all_branches if search_query.lower() in b.lower()
        ]

        if not filtered_branches:
            console.print(f"[red]No branches found matching '{search_query}'[/red]")
            # Default to main/master
            return "main" if "main" in all_branches else "master"
        
        selected = questionary.select(
            f"Select branch for {repo_name} (found {len(filtered_branches)} matches):",
            choices=filtered_branches[:50],  # Limit to 50 results
            default=filtered_branches[0],
        ).ask()

        if not selected:
            console.print("[red]Cancelled.[/red]")
            sys.exit(0)
    
    return selected


def main():
    console.print(
        Panel.fit("🌳 Git Worktree Creator with Submodule Support", style="bold blue")
    )

    # Check if we're in a git repository
    try:
        current_branch = get_current_branch()
    except subprocess.CalledProcessError:
        console.print("[red]Error: Not in a git repository![/red]")
        sys.exit(1)

    submodules = get_submodules()

    # Prompt for worktree name
    default_name = f"{current_branch}-worktree"
    # Clean up the name if it's from a detached state
    if default_name.startswith("detached-"):
        default_name = "worktree"
    
    worktree_name = questionary.text(
        "Enter worktree name:", 
        default=default_name
    ).ask()

    if not worktree_name:
        console.print("[red]Cancelled.[/red]")
        sys.exit(0)

    # Select main repository branch
    console.print("\n[bold]Main Repository Configuration:[/bold]")
    all_branches = get_remote_branches()
    main_branch = select_branch_interactive("main repository", all_branches, current_branch)

    # Select branches for each submodule
    submodule_branches = {}
    if submodules:
        console.print("\n[bold]Submodule Configuration:[/bold]")
        
        for submodule_path, current_sub_branch in submodules.items():
            console.print(f"\n[cyan]Submodule: {submodule_path}[/cyan]")
            
            # Get available branches for this submodule
            all_sub_branches = get_remote_branches(Path(submodule_path))
            selected_branch = select_branch_interactive(
                submodule_path, 
                all_sub_branches, 
                current_sub_branch
            )
            submodule_branches[submodule_path] = selected_branch

    # Create summary table
    table = Table(title="Worktree Configuration Summary")
    table.add_column("Repository", style="cyan")
    table.add_column("Branch", style="green")

    table.add_row("Main Repository", main_branch)
    for submodule, branch in submodule_branches.items():
        table.add_row(f"  └─ {submodule}", branch)

    console.print("\n")
    console.print(table)

    # Confirm
    if not questionary.confirm(
        "\nCreate worktree with this configuration?", default=True
    ).ask():
        console.print("[red]Cancelled.[/red]")
        sys.exit(0)

    # Create worktree
    current_dir = Path.cwd()
    parent_dir = current_dir.parent
    worktree_path = parent_dir / worktree_name

    console.print(f"\n[bold]Creating worktree at: {worktree_path}[/bold]")

    # Check if worktree already exists
    if worktree_path.exists():
        console.print(f"[red]Error: Directory {worktree_path} already exists![/red]")
        sys.exit(1)

    # Check if branch is already checked out
    worktree_list = run_command(["git", "worktree", "list", "--porcelain"], check=False)
    branch_already_checked_out = False
    
    if worktree_list.returncode == 0:
        for line in worktree_list.stdout.strip().split("\n"):
            if line.startswith("branch refs/heads/") and line.endswith(f"/{main_branch}"):
                branch_already_checked_out = True
                break

    if branch_already_checked_out:
        console.print(
            f"\n[yellow]Warning: Branch '{main_branch}' is already checked out in another worktree.[/yellow]"
        )
        # Create detached worktree and then switch to branch
        result = run_command(
            ["git", "worktree", "add", "--detach", str(worktree_path)],
            check=False,
        )
        if result.returncode == 0:
            # Switch to the desired branch
            run_command(
                ["git", "checkout", "-B", main_branch, f"origin/{main_branch}"], 
                cwd=worktree_path, 
                check=False
            )
    else:
        # Normal worktree creation
        result = run_command(
            ["git", "worktree", "add", str(worktree_path), main_branch], 
            check=False
        )

    if result.returncode != 0:
        console.print(f"[red]Error creating worktree: {result.stderr}[/red]")
        sys.exit(1)

    console.print("[green]✓ Worktree created successfully[/green]")

    # Initialize and update submodules
    if submodules:
        console.print("\n[bold]Initializing submodules...[/bold]")

        # Initialize submodules
        result = run_command(["git", "submodule", "init"], cwd=worktree_path)
        if result.returncode != 0:
            console.print(f"[red]Error initializing submodules: {result.stderr}[/red]")
        else:
            console.print("[green]✓ Submodules initialized[/green]")

        # Update each submodule to the specified branch
        for submodule_path, branch in submodule_branches.items():
            console.print(f"\nConfiguring {submodule_path} to use branch: {branch}")

            # Update submodule
            result = run_command(
                ["git", "submodule", "update", "--init", "--recursive", submodule_path],
                cwd=worktree_path,
                check=False,
            )

            if result.returncode != 0:
                console.print(
                    f"[yellow]Warning: Could not update {submodule_path}: {result.stderr}[/yellow]"
                )
                continue

            # Checkout the specified branch in the submodule
            submodule_full_path = worktree_path / submodule_path
            
            # Fetch latest changes
            run_command(["git", "fetch", "origin"], cwd=submodule_full_path, check=False)
            
            # Try to checkout the branch
            result = run_command(
                ["git", "checkout", branch], 
                cwd=submodule_full_path, 
                check=False
            )

            if result.returncode != 0:
                # Try to create and checkout branch from origin
                result = run_command(
                    ["git", "checkout", "-b", branch, f"origin/{branch}"],
                    cwd=submodule_full_path,
                    check=False,
                )

                if result.returncode != 0:
                    console.print(
                        f"[yellow]Warning: Could not checkout branch {branch} in {submodule_path}[/yellow]"
                    )
            else:
                console.print(f"[green]✓ {submodule_path} checked out to {branch}[/green]")

    # Copy environment files
    console.print("\n[bold]Copying environment files...[/bold]")
    copy_environment_files(current_dir, worktree_path)

    # Final summary
    console.print("\n" + "=" * 50)
    console.print("[bold green]✨ Worktree created successfully![/bold green]")
    console.print(f"\n[bold]Location:[/bold] {worktree_path.absolute()}")
    console.print("\n[bold]Next steps:[/bold]")
    console.print(f"  cd {worktree_path}")
    console.print(f"  # Start developing in your isolated worktree")
    console.print("\n[bold]To remove this worktree later:[/bold]")
    console.print(f"  git worktree remove {worktree_path}")
    console.print("\n[bold]To list all worktrees:[/bold]")
    console.print(f"  git worktree list")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        console.print("\n[red]Interrupted by user.[/red]")
        sys.exit(1)
    except Exception as e:
        console.print(f"\n[red]Unexpected error: {e}[/red]")
        sys.exit(1)

I want to create a worktree for a new branch where I have removed a previous dependent git submodule. But the script is asking me which branch I want to checkout on that submodule which should not exist in the new worktree.