afdo_tools: move worktree creation to git_utils

This will be needed soon by more CLs for b/332589934.

BUG=b:332589934
TEST=repo upload

Change-Id: Id91a3604e0e547de0d348e6aa01ec168e0165c50
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/5435854
Reviewed-by: Jordan Abrahams-Whitehead <ajordanr@google.com>
Commit-Queue: George Burgess <gbiv@chromium.org>
Tested-by: George Burgess <gbiv@chromium.org>

2 files changed, 46 insertions, 46 deletions

diff --git a/afdo_tools/update_kernel_afdo.py b/afdo_tools/update_kernel_afdo.py
index a7e40763..0a299bd2 100755
--- a/afdo_tools/update_kernel_afdo.py
+++ b/afdo_tools/update_kernel_afdo.py
@@ -9,7 +9,6 @@ It supports updating on canary, stable, and beta branches.
 """
 
 import argparse
-import contextlib
 import dataclasses
 import datetime
 import enum
@@ -21,8 +20,7 @@ import re
 import shlex
 import subprocess
 import sys
-import tempfile
-from typing import Dict, Generator, Iterable, List, Optional, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from cros_utils import git_utils
 
@@ -167,47 +165,6 @@ def get_parser():
     return parser
 
 
-@contextlib.contextmanager
-def git_worktree(git_directory: Path) -> Generator[Path, None, None]:
-    """Creates a temp worktree of `git_directory`, yielding the result."""
-    with tempfile.TemporaryDirectory(prefix="update_kernel_afdo_") as t:
-        tempdir = Path(t)
-        logging.info(
-            "Establishing worktree of %s in %s", git_directory, tempdir
-        )
-        subprocess.run(
-            [
-                "git",
-                "worktree",
-                "add",
-                "--detach",
-                "--force",
-                tempdir,
-            ],
-            cwd=git_directory,
-            check=True,
-            stdin=subprocess.DEVNULL,
-        )
-
-        try:
-            yield tempdir
-        finally:
-            # Explicitly `git worktree remove` here, so the parent worktree's
-            # metadata is cleaned up promptly.
-            subprocess.run(
-                [
-                    "git",
-                    "worktree",
-                    "remove",
-                    "--force",
-                    tempdir,
-                ],
-                cwd=git_directory,
-                check=False,
-                stdin=subprocess.DEVNULL,
-            )
-
-
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class GitBranch:
     """Represents a ChromeOS branch."""
@@ -807,7 +764,7 @@ def main(argv: List[str]) -> None:
 
     fetcher = KernelProfileFetcher()
     had_failures = False
-    with git_worktree(toolchain_utils) as worktree:
+    with git_utils.create_worktree(toolchain_utils) as worktree:
         for channel in opts.channel:
             branch = branches[channel]
             result = update_afdo_for_channel(
diff --git a/cros_utils/git_utils.py b/cros_utils/git_utils.py
index 1ae02ae0..6910a0ca 100644
--- a/cros_utils/git_utils.py
+++ b/cros_utils/git_utils.py
@@ -4,12 +4,14 @@
 
 """Shared utilities for working with git."""
 
+import contextlib
 import logging
 from pathlib import Path
 import re
 import shlex
 import subprocess
-from typing import Iterable, List
+import tempfile
+from typing import Generator, Iterable, List
 
 
 # Email address used to tag the detective as a reviewer.
@@ -105,3 +107,44 @@ def try_set_autosubmit_labels(chromeos_tree: Path, cl_id: int) -> None:
                 "Failed to run gerrit command %s. Ignoring.",
                 shlex.join(cmd),
             )
+
+
+@contextlib.contextmanager
+def create_worktree(git_directory: Path) -> Generator[Path, None, None]:
+    """Creates a temp worktree of `git_directory`, yielding the result."""
+    with tempfile.TemporaryDirectory(prefix="update_kernel_afdo_") as t:
+        tempdir = Path(t)
+        logging.info(
+            "Establishing worktree of %s in %s", git_directory, tempdir
+        )
+        subprocess.run(
+            [
+                "git",
+                "worktree",
+                "add",
+                "--detach",
+                "--force",
+                tempdir,
+            ],
+            cwd=git_directory,
+            check=True,
+            stdin=subprocess.DEVNULL,
+        )
+
+        try:
+            yield tempdir
+        finally:
+            # Explicitly `git worktree remove` here, so the parent worktree's
+            # metadata is cleaned up promptly.
+            subprocess.run(
+                [
+                    "git",
+                    "worktree",
+                    "remove",
+                    "--force",
+                    tempdir,
+                ],
+                cwd=git_directory,
+                check=False,
+                stdin=subprocess.DEVNULL,
+            )