Add temporary_worktree context manager

Add a context manager for creating temporary worktrees that are
automatically cleaned up on exit. This simplifies the common pattern
of creating a worktree for temporary operations.

Also add exist_ok parameter to add_worktree to support creating
worktrees in directories that already exist (e.g. from tempfile.mkdtemp).

NEWS

@@ -3,6 +3,12 @@
  * Fix worktree CLI tests to properly change to repository directory.
    (Jelmer Vernooij, #1738)
 
+ * Add ``temporary_worktree`` context manager for creating temporary worktrees
+   that are automatically cleaned up. (Jelmer Vernooij)
+
+ * Add ``exist_ok`` parameter to ``add_worktree`` to allow creation with
+   existing directories. (Jelmer Vernooij)
+
 0.24.1	2025-08-01
 
  * Require ``typing_extensions`` on Python 3.10.

dulwich/worktree.py

@@ -28,9 +28,12 @@ import os
 import shutil
 import stat
 import sys
+import tempfile
 import time
 import warnings
 from collections.abc import Iterable
+from contextlib import contextmanager
+from pathlib import Path
 
 from .errors import CommitError, HookError
 from .objects import Commit, ObjectID, Tag, Tree
@@ -140,6 +143,7 @@ class WorkTreeContainer:
         commit: ObjectID | None = None,
         force: bool = False,
         detach: bool = False,
+        exist_ok: bool = False,
     ) -> Repo:
         """Add a new worktree.
 
@@ -149,12 +153,19 @@ class WorkTreeContainer:
             commit: Specific commit to checkout (results in detached HEAD)
             force: Force creation even if branch is already checked out elsewhere
             detach: Detach HEAD in the new worktree
+            exist_ok: If True, do not raise an error if the directory already exists
 
         Returns:
             The newly created worktree repository
         """
         return add_worktree(
-            self._repo, path, branch=branch, commit=commit, force=force, detach=detach
+            self._repo,
+            path,
+            branch=branch,
+            commit=commit,
+            force=force,
+            detach=detach,
+            exist_ok=exist_ok,
         )
 
     def remove(self, path: str | bytes | os.PathLike, force: bool = False) -> None:
@@ -806,6 +817,7 @@ def add_worktree(
     commit: ObjectID | None = None,
     force: bool = False,
     detach: bool = False,
+    exist_ok: bool = False,
 ) -> Repo:
     """Add a new worktree to the repository.
 
@@ -816,12 +828,13 @@ def add_worktree(
         commit: Specific commit to checkout (results in detached HEAD)
         force: Force creation even if branch is already checked out elsewhere
         detach: Detach HEAD in the new worktree
+        exist_ok: If True, do not raise an error if the directory already exists
 
     Returns:
         The newly created worktree repository
 
     Raises:
-        ValueError: If the path already exists or branch is already checked out
+        ValueError: If the path already exists (and exist_ok is False) or branch is already checked out
     """
     from .repo import Repo as RepoClass
 
@@ -830,7 +843,7 @@ def add_worktree(
         path = os.fsdecode(path)
 
     # Check if path already exists
-    if os.path.exists(path):
+    if os.path.exists(path) and not exist_ok:
         raise ValueError(f"Path already exists: {path}")
 
     # Normalize branch name
@@ -871,7 +884,7 @@ def add_worktree(
         detach = True
 
     # Create the worktree directory
-    os.makedirs(path)
+    os.makedirs(path, exist_ok=exist_ok)
 
     # Initialize the worktree
     identifier = os.path.basename(path)
@@ -1141,3 +1154,36 @@ def move_worktree(
     # Update the gitdir pointer in the control directory
     with open(os.path.join(worktree_control_dir, GITDIR), "wb") as f:
         f.write(os.fsencode(gitdir_file) + b"\n")
+
+
+@contextmanager
+def temporary_worktree(repo, prefix="tmp-worktree-"):
+    """Create a temporary worktree that is automatically cleaned up.
+
+    Args:
+        repo: Dulwich repository object
+        prefix: Prefix for the temporary directory name
+
+    Yields:
+        Worktree object
+    """
+    temp_dir = None
+    worktree = None
+
+    try:
+        # Create temporary directory
+        temp_dir = tempfile.mkdtemp(prefix=prefix)
+
+        # Add worktree
+        worktree = repo.worktrees.add(temp_dir, exist_ok=True)
+
+        yield worktree
+
+    finally:
+        # Clean up worktree registration
+        if worktree:
+            repo.worktrees.remove(worktree.path)
+
+        # Clean up temporary directory
+        if temp_dir and Path(temp_dir).exists():
+            shutil.rmtree(temp_dir)

tests/test_worktree.py

@@ -39,6 +39,7 @@ from dulwich.worktree import (
     move_worktree,
     prune_worktrees,
     remove_worktree,
+    temporary_worktree,
     unlock_worktree,
 )
 
@@ -719,3 +720,109 @@ class WorkTreeOperationsTests(WorkTreeTestCase):
         with self.assertRaises(ValueError) as cm:
             move_worktree(self.repo, wt_path, new_path)
         self.assertIn("Path already exists", str(cm.exception))
+
+
+class TemporaryWorktreeTests(TestCase):
+    """Tests for temporary_worktree context manager."""
+
+    def setUp(self) -> None:
+        super().setUp()
+        self.tempdir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.tempdir)
+        self.repo_path = os.path.join(self.tempdir, "repo")
+        self.repo = Repo.init(self.repo_path, mkdir=True)
+
+        # Create an initial commit so HEAD exists
+        readme_path = os.path.join(self.repo_path, "README.md")
+        with open(readme_path, "w") as f:
+            f.write("# Test Repository\n")
+        porcelain.add(self.repo, [readme_path])
+        porcelain.commit(self.repo, message=b"Initial commit")
+
+    def test_temporary_worktree_creates_and_cleans_up(self) -> None:
+        """Test that temporary worktree is created and cleaned up."""
+        worktree_path = None
+
+        # Use the context manager
+        with temporary_worktree(self.repo) as worktree:
+            worktree_path = worktree.path
+
+            # Check that worktree exists
+            self.assertTrue(os.path.exists(worktree_path))
+
+            # Check that it's in the list of worktrees
+            worktrees = list_worktrees(self.repo)
+            paths = [wt.path for wt in worktrees]
+            self.assertIn(worktree_path, paths)
+
+            # Check that .git file exists in worktree
+            gitdir_file = os.path.join(worktree_path, ".git")
+            self.assertTrue(os.path.exists(gitdir_file))
+
+        # After context manager exits, check cleanup
+        self.assertFalse(os.path.exists(worktree_path))
+
+        # Check that it's no longer in the list of worktrees
+        worktrees = list_worktrees(self.repo)
+        paths = [wt.path for wt in worktrees]
+        self.assertNotIn(worktree_path, paths)
+
+    def test_temporary_worktree_with_custom_prefix(self) -> None:
+        """Test temporary worktree with custom prefix."""
+        custom_prefix = "my-custom-prefix-"
+
+        with temporary_worktree(self.repo, prefix=custom_prefix) as worktree:
+            # Check that the directory name starts with our prefix
+            dirname = os.path.basename(worktree.path)
+            self.assertTrue(dirname.startswith(custom_prefix))
+
+    def test_temporary_worktree_cleanup_on_exception(self) -> None:
+        """Test that cleanup happens even when exception is raised."""
+        worktree_path = None
+
+        class TestException(Exception):
+            pass
+
+        try:
+            with temporary_worktree(self.repo) as worktree:
+                worktree_path = worktree.path
+                self.assertTrue(os.path.exists(worktree_path))
+                raise TestException("Test exception")
+        except TestException:
+            pass
+
+        # Cleanup should still happen
+        self.assertFalse(os.path.exists(worktree_path))
+
+        # Check that it's no longer in the list of worktrees
+        worktrees = list_worktrees(self.repo)
+        paths = [wt.path for wt in worktrees]
+        self.assertNotIn(worktree_path, paths)
+
+    def test_temporary_worktree_operations(self) -> None:
+        """Test that operations can be performed in temporary worktree."""
+        # Create a test file in main repo
+        test_file = os.path.join(self.repo_path, "test.txt")
+        with open(test_file, "w") as f:
+            f.write("Hello, world!")
+
+        porcelain.add(self.repo, [test_file])
+        porcelain.commit(self.repo, message=b"Initial commit")
+
+        with temporary_worktree(self.repo) as worktree:
+            # Check that the file exists in the worktree
+            wt_test_file = os.path.join(worktree.path, "test.txt")
+            self.assertTrue(os.path.exists(wt_test_file))
+
+            # Read and verify content
+            with open(wt_test_file) as f:
+                content = f.read()
+            self.assertEqual(content, "Hello, world!")
+
+            # Make changes in the worktree
+            with open(wt_test_file, "w") as f:
+                f.write("Modified content")
+
+            # Changes should be visible in status
+            status = porcelain.status(worktree)
+            self.assertIn(b"test.txt", status.unstaged)