gitworktree/gitworktree/core.py at main · gitstq/gitworktree

672 lines (562 loc) · 18.2 KB
Core worktree operations.
High-level functions that combine git operations with display and validation logic.
import shutil
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from .config import (
    DEFAULT_FORMAT,
    DEFAULT_PATH_TEMPLATE,
    DEFAULT_SORT,
    FILTER_KEYS,
    FORMAT_KEYS,
    SORT_KEYS,
    STATUS_CLEAN,
    STATUS_DIRTY,
    WORKTREE_BASE_DIR,
from .display import (
    format_worktree_compact,
    format_worktree_info,
    format_worktree_json,
    format_worktree_table,
    interactive_picker,
    print_error,
    print_success,
    print_warning,
from .git import GitCommandError, GitRepo
from .utils import (
    confirm,
    find_repo_root,
    get_worktree_path,
    is_valid_branch_name,
    parse_duration,
    sanitize_branch_name,
    truncate_string,
def cmd_list(
    repo: GitRepo,
    fmt: str = DEFAULT_FORMAT,
    sort: str = DEFAULT_SORT,
    filter_status: Optional[str] = None,
    style: Optional[Style] = None,
    List all worktrees with formatted output.
        repo: GitRepo instance.
        fmt: Output format (table/json/compact).
        sort: Sort key (name/date/branch).
        filter_status: Filter by status (clean/dirty).
        style: Style instance for colorized output.
    Returns:
        Formatted output string.
    s = style or Style()
    if fmt not in FORMAT_KEYS:
        print_error(f"Invalid format: {fmt}. Choose from: {', '.join(FORMAT_KEYS)}", s)
        return ""
    if sort not in SORT_KEYS:
        print_error(f"Invalid sort key: {sort}. Choose from: {', '.join(SORT_KEYS)}", s)
        return ""
    if filter_status and filter_status not in FILTER_KEYS:
        print_error(f"Invalid filter: {filter_status}. Choose from: {', '.join(FILTER_KEYS)}", s)
        return ""
    # Get raw worktree list
    raw_worktrees = repo.list_worktrees()
    if not raw_worktrees:
        return s.yellow("No worktrees found.")
    # Enrich with status information
    worktrees = []
    for wt in raw_worktrees:
        branch = wt.get("branch", "unknown")
        enriched = dict(wt)
        # Get branch status
        if branch != "detached":
            try:
                status_info = repo.get_branch_status(branch)
                enriched["status"] = status_info["status"]
                enriched["ahead"] = status_info["ahead"]
                enriched["behind"] = status_info["behind"]
                enriched["dirty"] = status_info["dirty"]
            except GitCommandError:
                enriched["status"] = STATUS_CLEAN
                enriched["ahead"] = 0
                enriched["behind"] = 0
                enriched["dirty"] = False
        else:
            enriched["status"] = STATUS_CLEAN
            enriched["ahead"] = 0
            enriched["behind"] = 0
            enriched["dirty"] = False
        # Get last modified time
        try:
            modified = repo.get_last_modified(wt["path"])
            enriched["last_modified"] = _format_relative_time(modified)
        except GitCommandError:
            enriched["last_modified"] = ""
        worktrees.append(enriched)
    # Apply filter
    if filter_status:
        worktrees = [wt for wt in worktrees if wt.get("status") == filter_status]
    if sort == "name":
        worktrees.sort(key=lambda w: os.path.basename(w.get("path", "")))
    elif sort == "branch":
        worktrees.sort(key=lambda w: w.get("branch", ""))
    elif sort == "date":
        worktrees.sort(key=lambda w: w.get("last_modified", ""), reverse=True)
    # Format output
    if fmt == "json":
        return format_worktree_json(worktrees)
    elif fmt == "compact":
        return format_worktree_compact(worktrees, s)
        return format_worktree_table(worktrees, s)
def cmd_add(
    repo: GitRepo,
    branch: str,
    create_branch: bool = False,
    track: Optional[str] = None,
    path_template: str = DEFAULT_PATH_TEMPLATE,
    base_dir: str = WORKTREE_BASE_DIR,
    style: Optional[Style] = None,
    Create a new worktree for a branch.
        repo: GitRepo instance.
        branch: Branch name to check out.
        create_branch: If True, create a new branch.
        track: Remote branch to track.
        path_template: Template for the worktree directory name.
        base_dir: Base directory for worktrees.
        style: Style instance for colorized output.
    Returns:
        True if the worktree was created successfully.
    s = style or Style()
    # Validate branch name
    if create_branch and not is_valid_branch_name(branch):
        print_error(f"Invalid branch name: {branch}", s)
        return False
    # Check if branch already has a worktree
    existing = repo.list_worktrees()
    for wt in existing:
        if wt.get("branch") == branch:
            print_error(
                f"Branch '{s.green(branch)}' already has a worktree at: {s.cyan(wt['path'])}",
            return False
    # Get commit info for template variables
        commit_info = repo.get_commit_info(branch)
        short_hash = commit_info.get("short_hash", "")
        author = commit_info.get("author", "")
    except GitCommandError:
        short_hash = ""
        author = ""
    # Generate worktree path
    wt_path = get_worktree_path(
        repo_root=repo.get_root(),
        branch=branch,
        template=path_template,
        base_dir=base_dir,
        short_hash=short_hash,
        author=author,
    # Check if path already exists
    if wt_path.exists():
        print_error(
            f"Directory already exists: {s.cyan(str(wt_path))}",
        return False
    # Create the worktree
        track_arg = track
        if track and not track_arg.startswith("origin/"):
            track_arg = f"origin/{track}"
        repo.add_worktree(
            path=str(wt_path),
            branch=branch if not create_branch else None,
            create_branch=create_branch,
            track=track_arg,
        print_success(
            f"Worktree created: {s.cyan(str(wt_path))} -> {s.green(branch)}",
        return True
    except GitCommandError as e:
        print_error(str(e), s)
        return False
def cmd_switch(
    repo: GitRepo,
    shell: bool = False,
    style: Optional[Style] = None,
) -> Optional[str]:
    Interactive worktree switcher.
        repo: GitRepo instance.
        shell: If True, output eval-compatible shell command.
        style: Style instance for colorized output.
    Returns:
        Shell command string if shell mode, None otherwise.
    s = style or Style()
    worktrees = repo.list_worktrees()
    if not worktrees:
        print_error("No worktrees found.", s)
        return None
    # Build picker items
    items: List[Tuple[str, str]] = []
    for wt in worktrees:
        branch = wt.get("branch", "unknown")
        path = wt.get("path", "")
        is_main = wt.get("is_main", False)
        label = f"{branch}" + (" (main)" if is_main else "")
        items.append((label, path))
    # Show picker
    index = interactive_picker(items, s)
    if index < 0:
        print_warning("Selection cancelled.", s)
        return None
    selected_path = worktrees[index].get("path", "")
    selected_branch = worktrees[index].get("branch", "unknown")
    if shell:
        return f"cd {selected_path}"
        print(f"\n{s.bold('Switch to worktree:')}")
        print(f"  Branch: {s.green(selected_branch)}")
        print(f"  Path:   {s.cyan(selected_path)}")
        print(f"\nRun: {s.bold(f'cd {selected_path}')}")
        return None
def cmd_remove(
    repo: GitRepo,
    identifier: str,
    force: bool = False,
    prune: bool = True,
    style: Optional[Style] = None,
    Remove a worktree by branch name or path.
        repo: GitRepo instance.
        identifier: Branch name or worktree path.
        force: If True, skip confirmation.
        prune: If True, prune stale entries after removal.
        style: Style instance for colorized output.
    Returns:
        True if the worktree was removed successfully.
    s = style or Style()
    # Find the worktree
    worktrees = repo.list_worktrees()
    target_wt = None
    for wt in worktrees:
        if wt.get("branch") == identifier:
            target_wt = wt
            break
        if os.path.abspath(wt.get("path", "")) == os.path.abspath(identifier):
            target_wt = wt
            break
    if not target_wt:
        print_error(f"No worktree found for: {identifier}", s)
        return False
    # Don't allow removing the main worktree
    if target_wt.get("is_main", False):
        print_error("Cannot remove the main worktree.", s)
        return False
    branch = target_wt.get("branch", "unknown")
    wt_path = target_wt.get("path", "")
    # Check if branch is merged
    merged = False
        merged = repo.is_branch_merged(branch)
    except GitCommandError:
        pass
    # Confirmation
    if not force:
        msg = f"Remove worktree '{s.green(branch)}' at {s.cyan(wt_path)}?"
        if merged:
            msg += f" {s.yellow('(branch is merged)')}"
        if not confirm(msg, default=False):
            print_warning("Removal cancelled.", s)
            return False
    # Remove the worktree
        repo.remove_worktree(wt_path, force=force)
        print_success(f"Worktree removed: {s.green(branch)}", s)
        # Prune stale entries
        if prune:
            try:
                repo.prune_worktrees()
            except GitCommandError:
        # Suggest branch deletion
        if merged:
            delete_msg = f"Branch '{branch}' is merged. "
            if force:
                    repo.delete_branch(branch)
                    print_success(f"Branch deleted: {s.green(branch)}", s)
                except GitCommandError as e:
                    print_warning(f"Could not delete branch: {e}", s)
            else:
                print_warning(
                    f"Branch '{branch}' is merged. "
                    f"Consider deleting it with: git branch -d {branch}",
        return True
    except GitCommandError as e:
        print_error(str(e), s)
        return False
def cmd_merge(
    repo: GitRepo,
    source: str,
    target: Optional[str] = None,
    delete: bool = False,
    no_ff: bool = False,
    style: Optional[Style] = None,
    Merge a worktree branch into a target branch.
        repo: GitRepo instance.
        source: Source branch to merge from.
        target: Target branch. Defaults to default branch.
        delete: If True, delete the source branch and worktree after merge.
        no_ff: If True, use --no-ff merge strategy.
        style: Style instance for colorized output.
    Returns:
        True if the merge was successful.
    s = style or Style()
    if target is None:
        target = repo.get_default_branch()
    # Verify source branch exists
    branches = repo.get_branches()
    if source not in branches:
        print_error(f"Branch '{source}' not found.", s)
        return False
    # Verify target branch exists
    if target not in branches:
        print_error(f"Target branch '{target}' not found.", s)
        return False
    # Check if already merged
    if repo.is_branch_merged(source, target):
        print_warning(
            f"Branch '{source}' is already merged into '{target}'.",
        if not confirm("Continue anyway?", default=False):
            return False
    # Perform merge
        print(f"Merging {s.green(source)} into {s.green(target)}...")
        repo.merge_branch(source, target, no_ff=no_ff)
        print_success(f"Successfully merged '{source}' into '{target}'.", s)
        # Optionally delete branch and worktree
        if delete:
            # Find and remove worktree
            worktrees = repo.list_worktrees()
            for wt in worktrees:
                if wt.get("branch") == source and not wt.get("is_main", False):
                        repo.remove_worktree(wt["path"], force=True)
                        print_success(f"Worktree removed: {source}", s)
                    except GitCommandError as e:
                        print_warning(f"Could not remove worktree: {e}", s)
            # Delete branch
            try:
                repo.delete_branch(source, force=True)
                print_success(f"Branch deleted: {source}", s)
            except GitCommandError as e:
                print_warning(f"Could not delete branch: {e}", s)
            # Prune
            try:
                repo.prune_worktrees()
            except GitCommandError:
        return True
    except GitCommandError as e:
        print_error(f"Merge failed: {e}", s)
        return False
def cmd_info(
    repo: GitRepo,
    identifier: str,
    style: Optional[Style] = None,
    Show detailed information about a worktree.
        repo: GitRepo instance.
        identifier: Branch name or worktree path.
        style: Style instance for colorized output.
    Returns:
        Formatted information string.
    s = style or Style()
    # Find the worktree
    worktrees = repo.list_worktrees()
    target_wt = None
    for wt in worktrees:
        if wt.get("branch") == identifier:
            target_wt = wt
            break
        if os.path.abspath(wt.get("path", "")) == os.path.abspath(identifier):
            target_wt = wt
            break
    if not target_wt:
        print_error(f"No worktree found for: {identifier}", s)
        return ""
    branch = target_wt.get("branch", "unknown")
    wt_path = target_wt.get("path", "")
    # Get commit info
        commit_info = repo.get_commit_info("HEAD", worktree_path=wt_path)
    except GitCommandError:
        commit_info = {}
    # Get branch status
    status_info = {}
    if branch != "detached":
        try:
            status_info = repo.get_branch_status(branch)
        except GitCommandError:
            status_info = {"status": "clean", "ahead": 0, "behind": 0}
    # Get uncommitted changes
        changes = repo.get_uncommitted_changes(wt_path)
    except GitCommandError:
        changes = {}
    # Build info dict
    info: Dict[str, Any] = {
        "branch": branch,
        "path": wt_path,
        "commit": commit_info.get("short_hash", target_wt.get("commit", "")[:7]),
        "author": commit_info.get("author", ""),
        "date": commit_info.get("date", ""),
        "status": status_info.get("status", "clean"),
        "ahead": status_info.get("ahead", 0),
        "behind": status_info.get("behind", 0),
        "changes": changes,
        "is_main": target_wt.get("is_main", False),
    return format_worktree_info(info, s)
def cmd_clean(
    repo: GitRepo,
    dry_run: bool = False,
    force: bool = False,
    style: Optional[Style] = None,
    Clean orphaned worktrees whose branches have been deleted.
        repo: GitRepo instance.
        dry_run: If True, only show what would be cleaned.
        force: If True, skip confirmation.
        style: Style instance for colorized output.
    Returns:
        True if cleanup was successful.
    s = style or Style()
    worktrees = repo.list_worktrees()
    branches = repo.get_branches()
    orphaned: List[Dict[str, Any]] = []
    for wt in worktrees:
        if wt.get("is_main", False):
            continue
        branch = wt.get("branch", "")
        if branch == "detached":
            continue
        if branch not in branches:
            orphaned.append(wt)
    if not orphaned:
        print_success("No orphaned worktrees found.", s)
        return True
    # Show orphaned worktrees
    print(s.bold(f"Found {len(orphaned)} orphaned worktree(s):"))
    for wt in orphaned:
        branch = wt.get("branch", "unknown")
        path = wt.get("path", "")
        print(f"  {s.red(branch)}  {s.cyan(path)}")
    if dry_run:
        print(s.yellow("\nDry run: no changes made."))
        return True
    # Confirmation
    if not force:
        if not confirm(f"Remove {len(orphaned)} orphaned worktree(s)?", default=False):
            print_warning("Cleanup cancelled.", s)
            return False
    # Remove orphaned worktrees
    success = True
    for wt in orphaned:
        branch = wt.get("branch", "unknown")
        path = wt.get("path", "")
        try:
            repo.remove_worktree(path, force=True)
            print_success(f"Removed: {branch}", s)
        except GitCommandError as e:
            print_error(f"Failed to remove {branch}: {e}", s)
            success = False
    # Prune
        repo.prune_worktrees()
    except GitCommandError:
        pass
    if success:
        print_success(f"Cleaned {len(orphaned)} orphaned worktree(s).", s)
    return success
def cmd_path(
    repo: GitRepo,
    branch: str,
    style: Optional[Style] = None,
) -> Optional[str]:
    Get the worktree path for a given branch.
        repo: GitRepo instance.
        branch: Branch name to look up.
        style: Style instance for colorized output.
    Returns:
        Worktree path string, or None if not found.
    s = style or Style()
    worktrees = repo.list_worktrees()
    for wt in worktrees:
        if wt.get("branch") == branch:
            return wt.get("path")
    print_error(f"No worktree found for branch: {branch}", s)
    return None
def _format_relative_time(date_str: str) -> str:
    Format an ISO date string as a relative time string.
        date_str: ISO 8601 date string.
    Returns:
        Relative time string (e.g., "2h ago", "3d ago").
    if not date_str:
        return ""
        # Parse ISO 8601 format
        dt = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
        now = datetime.now(dt.tzinfo) if dt.tzinfo else datetime.now()
        diff = now - dt
        seconds = int(diff.total_seconds())
        if seconds < 0:
            return "just now"
        return parse_duration(seconds) + " ago"
    except (ValueError, TypeError):
        return date_str
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

core.py

Latest commit

History

core.py

File metadata and controls
