From dab87a6d81c1b419fcfc4ce53f6881892b0fcfe0 Mon Sep 17 00:00:00 2001 From: Ankit Khullar Date: Sun, 18 Jan 2026 19:19:17 +0530 Subject: [PATCH] Refactor worktree management and enhance tests - Updated `add_worktree` function to include a `create_branch` parameter for better branch handling. - Modified tests to validate worktree addition with both new and existing branches. - Improved CLI messages for switching worktrees to provide clearer instructions. - Ensured existing worktree checks are more robust in the `WorktreeManager` class. These changes enhance the functionality and usability of the worktree management system. --- tools/wt-worktree/tests/test_git.py | 16 ++++++++++++++-- tools/wt-worktree/wt/cli.py | 13 +++---------- tools/wt-worktree/wt/git.py | 26 ++++++++++++++++++-------- tools/wt-worktree/wt/worktree.py | 22 ++++++++-------------- 4 files changed, 43 insertions(+), 34 deletions(-) diff --git a/tools/wt-worktree/tests/test_git.py b/tools/wt-worktree/tests/test_git.py index d90314a..0a3a5ea 100644 --- a/tools/wt-worktree/tests/test_git.py +++ b/tools/wt-worktree/tests/test_git.py @@ -68,7 +68,19 @@ def test_list_worktrees(git_repo): def test_add_worktree(git_repo, temp_dir): """Test adding a worktree.""" wt_path = temp_dir / "test-worktree" - git.add_worktree(wt_path, "test-branch", "HEAD", repo_path=git_repo) + git.add_worktree(wt_path, "test-branch", True, "HEAD", repo_path=git_repo) + + assert wt_path.exists() + assert git.branch_exists("test-branch", git_repo) + + worktrees = git.list_worktrees(git_repo) + assert len(worktrees) == 2 + +def test_add_worktree_with_existing_branch(git_repo, temp_dir): + """Test adding a worktree with an existing branch.""" + git.create_branch("test-branch", "HEAD", git_repo) + wt_path = temp_dir / "test-worktree" + git.add_worktree(wt_path, "test-branch", False, "HEAD", repo_path=git_repo) assert wt_path.exists() assert git.branch_exists("test-branch", git_repo) @@ -80,7 +92,7 @@ def test_add_worktree(git_repo, temp_dir): def test_remove_worktree(git_repo, temp_dir): """Test removing a worktree.""" wt_path = temp_dir / "test-worktree" - git.add_worktree(wt_path, "test-branch", "HEAD", repo_path=git_repo) + git.add_worktree(wt_path, "test-branch", True, "HEAD", repo_path=git_repo) git.remove_worktree(wt_path, repo_path=git_repo) assert not wt_path.exists() diff --git a/tools/wt-worktree/wt/cli.py b/tools/wt-worktree/wt/cli.py index 3e8d68c..4811a7a 100644 --- a/tools/wt-worktree/wt/cli.py +++ b/tools/wt-worktree/wt/cli.py @@ -78,7 +78,7 @@ def init(ctx: Context, prefix: str, path_pattern: str): success(f"Configuration saved to {config_path}") info(f"\nBranch prefix: {prefix}") info(f"Path pattern: {path_pattern}") - info(f"\nConfiguration will be used for all repositories.") + info("\nConfiguration will be used for all repositories.") except ConfigError as e: error(str(e), EXIT_ERROR) @@ -148,20 +148,13 @@ def switch(ctx: Context, name: Optional[str], create: bool, base: Optional[str], if shell_helper: print(target_wt["path"]) else: - success(f"Switched to worktree '{name}' at {target_wt['path']}") + success(f"To switch to worktree '{name}' run: cd {target_wt['path']}") elif create: # Create new worktree try: # Check if branch exists full_branch = ctx.config.get_branch_name(name) - if git.branch_exists(full_branch, ctx.repo_root): - if not confirm( - f"Branch '{full_branch}' already exists.\n" - "Create worktree for existing branch?", - default=True - ): - sys.exit(EXIT_CANCELLED) # Create worktree wt_path = ctx.manager.create_worktree(name, base, detached) @@ -175,7 +168,7 @@ def switch(ctx: Context, name: Optional[str], create: bool, base: Optional[str], if shell_helper: print(wt_path) else: - success(f"Created and switched to worktree '{name}' at {wt_path}") + success(f"To switch to worktree '{name}' run: cd {wt_path}") except git.GitError as e: error(str(e), EXIT_GIT_ERROR) diff --git a/tools/wt-worktree/wt/git.py b/tools/wt-worktree/wt/git.py index a0afec6..19e73c0 100644 --- a/tools/wt-worktree/wt/git.py +++ b/tools/wt-worktree/wt/git.py @@ -212,14 +212,19 @@ def worktree_exists(name: str, path: Optional[Path] = None) -> Tuple[bool, Optio return False, None -def add_worktree(path: Path, branch: str, base: Optional[str] = None, - detached: bool = False, repo_path: Optional[Path] = None): +def add_worktree(path: Path, + branch: str, + create_branch: bool = False, + base: Optional[str] = None, + detached: bool = False, + repo_path: Optional[Path] = None): """ Create a new worktree. Args: path: Path where worktree will be created branch: Branch name for the worktree + create_branch: create a new branch instead of using an existing one base: Base branch/commit (if None, uses current HEAD) detached: Create in detached HEAD state repo_path: Path to main repo (for running command) @@ -228,13 +233,18 @@ def add_worktree(path: Path, branch: str, base: Optional[str] = None, if detached: args.append("--detach") - else: + args.append(str(path)) + if base: + args.append(base) + elif create_branch: args.extend(["-b", branch]) - - args.append(str(path)) - - if base: - args.append(base) + args.append(str(path)) + if base: + args.append(base) + else: + # Use existing branch - format: git worktree add + args.append(str(path)) + args.append(branch) run_git(args, cwd=repo_path) diff --git a/tools/wt-worktree/wt/worktree.py b/tools/wt-worktree/wt/worktree.py index 62a2f46..7d87990 100644 --- a/tools/wt-worktree/wt/worktree.py +++ b/tools/wt-worktree/wt/worktree.py @@ -141,20 +141,14 @@ def create_worktree(self, name: str, base: Optional[str] = None, """ # Get full branch name branch = self.config.get_branch_name(name) + create_branch = not git.branch_exists(branch, self.repo_root) - # Check if branch already exists - if git.branch_exists(branch, self.repo_root): - exists, path = git.worktree_exists(branch, self.repo_root) - if exists: - raise git.GitError( - f"Worktree '{name}' already exists at {path}\n" - f"Use 'wt switch {name}' to switch to it." - ) - else: - raise git.GitError( - f"Branch '{branch}' already exists but has no worktree.\n" - f"Use 'git worktree add' manually or delete the branch first." - ) + exists, path = git.worktree_exists(branch, self.repo_root) + if exists: + raise git.GitError( + f"Worktree '{name}' already exists at {path}\n" + f"Use 'wt switch {name}' to switch to it." + ) # Resolve worktree path wt_path = self.config.resolve_path_pattern(name, branch) @@ -172,7 +166,7 @@ def create_worktree(self, name: str, base: Optional[str] = None, # Create worktree try: - git.add_worktree(wt_path, branch, base, detached, self.repo_root) + git.add_worktree(wt_path, branch, create_branch, base, detached, self.repo_root) except git.GitError as e: raise git.GitError(f"Failed to create worktree: {e}")