fix: CG-9440 clean repo - clears from the default branch

src/codegen/git/repo_operator/remote_repo_operator.py

@@ -24,7 +24,7 @@
     """A wrapper around GitPython to make it easier to interact with a cloned lowside repo."""
 
     # __init__ attributes
     repo_config: RepoConfig
     base_dir: str
     access_token: str | None = None
 
@@ -43,7 +43,7 @@
         bot_commit: bool = True,
         access_token: str | None = None,
     ) -> None:
         super().__init__(repo_config=repo_config, base_dir=base_dir, bot_commit=bot_commit)
         self.access_token = access_token
         self.setup_repo_dir(setup_option=setup_option, shallow=shallow)
 
@@ -79,6 +79,22 @@
     # SET UP
     ####################################################################################################################
 
+    @override
+    def clean_repo(self) -> None:
+        """Cleans the repo by:
+        1. Discards any changes (tracked/untracked)
+        2. Checks out the default branch (+ makes sure it's up to date with the remote)
+        3. Deletes all branches except the default branch
+        4. Deletes all remotes except origin
+
+        Used in SetupOption.PULL_OR_CLONE to allow people to re-use existing repos and start from a clean state.
+        """
+        logger.info(f"Cleaning repo at {self.repo_path} ...")
+        self.discard_changes()
+        self.checkout_branch(self.default_branch, remote=True)
+        self.clean_branches()
+        self.clean_remotes()
+
     @override
     def pull_repo(self) -> None:
         """Pull the latest commit down to an existing local repo"""

src/codegen/git/repo_operator/repo_operator.py

@@ -57,7 +57,7 @@
 
     @property
     def clone_url(self) -> str:
         return f"https://github.com/{self.repo_config.full_name}.git"
 
     @property
     def viz_path(self) -> str:
@@ -98,7 +98,7 @@
         email_level = None
         levels = ["system", "global", "user", "repository"]
         for level in levels:
             with git_cli.config_reader(level) as reader:
                 if reader.has_option("user", "name") and not username:
                     username = reader.get("user", "name")
                     user_level = level
@@ -159,15 +159,11 @@
     def clean_repo(self) -> None:
         """Cleans the repo by:
         1. Discards any changes (tracked/untracked)
-        2. Checks out the default branch (+ makes sure it's up to date with the remote)
-        3. Deletes all branches except the default branch
-        4. Deletes all remotes except origin
-
-        Used in SetupOption.PULL_OR_CLONE to allow people to re-use existing repos and start from a clean state.
+        2. Deletes all branches except the checked out branch
+        3. Deletes all remotes except origin
         """
         logger.info(f"Cleaning repo at {self.repo_path} ...")
         self.discard_changes()
-        self.checkout_branch(self.default_branch)  # TODO(CG-9440): add back remote=True
         self.clean_branches()
         self.clean_remotes()
 
@@ -277,20 +273,6 @@
             return False
         return self.git_cli.active_branch.name == branch_name
 
-    def delete_local_branch(self, branch_name: str) -> None:
-        if branch_name not in self.git_cli.branches:
-            logger.info(f"Branch {branch_name} does not exist locally. Skipping delete_local_branch.")
-            return
-        if branch_name is self.default_branch:
-            msg = "Deleting the default branch is not implemented yet."
-            raise NotImplementedError(msg)
-
-        if self.is_branch_checked_out(branch_name):
-            self.checkout_branch(self.default_branch)
-
-        logger.info(f"Deleting local branch: {branch_name} ...")
-        self.git_cli.delete_head(branch_name, force=True)  # force deletes even if the branch has unmerged changes
-
     def checkout_branch(self, branch_name: str | None, *, remote: bool = False, remote_name: str = "origin", create_if_missing: bool = True) -> CheckoutResult:
         """Attempts to check out the branch in the following order:
         - Check out the local branch by name
@@ -449,7 +431,7 @@
                     return content
             except UnicodeDecodeError:
                 print(f"Warning: Unable to decode file {file_path}. Skipping.")
                 return None
 
     def write_file(self, relpath: str, content: str) -> None:
         """Writes file content to disk"""
@@ -494,7 +476,7 @@
 
         # Iterate through files and yield contents
         for rel_filepath in filepaths:
             rel_filepath: str
             filepath = os.path.join(self.repo_path, rel_filepath)
 
             # Filter by subdirectory (includes full filenames)
@@ -525,7 +507,7 @@
         list_files = []
 
         for rel_filepath in self.git_cli.git.ls_files().split("\n"):
             rel_filepath: str
             if subdirs and not any(d in rel_filepath for d in subdirs):
                 continue
             if extensions is None or any(rel_filepath.endswith(e) for e in extensions):
@@ -549,7 +531,7 @@
 
     def get_modified_files_in_last_n_days(self, days: int = 1) -> tuple[list[str], list[str]]:
         """Returns a list of files modified and deleted in the last n days"""
         modified_files = []
         deleted_files = []
         allowed_extensions = [".py"]
 
@@ -565,9 +547,9 @@
                     if file in modified_files:
                         modified_files.remove(file)
                 else:
                     if file not in modified_files and file[-3:] in allowed_extensions:
                         modified_files.append(file)
         return modified_files, deleted_files
 
     @abstractmethod
     def base_url(self) -> str | None: ...

tests/integration/codegen/git/repo_operator/test_remote_repo_operator.py

@@ -12,7 +12,12 @@
 
 @pytest.fixture
 def op(repo_config, request, tmpdir):
-    op = RemoteRepoOperator(repo_config, shallow=request.param, base_dir=tmpdir, bot_commit=False)
+    op = RemoteRepoOperator(
+        repo_config,
+        shallow=request.param if hasattr(request, "param") else True,
+        base_dir=tmpdir,
+        bot_commit=False,
+    )
     yield op
 
 
@@ -76,3 +81,24 @@ def test_checkout_branch_remote_already_checked_out_resets_branch(mock_git_clien
     assert res == CheckoutResult.SUCCESS
     assert len(op.git_cli.heads) == 1
     assert op.head_commit.hexsha == original_commit_head.hexsha
+
+
+def test_clean_repo(op: RemoteRepoOperator):
+    num_branches = len(op.git_cli.branches)
+    op.checkout_branch(branch_name="test_branch", create_if_missing=True)
+    with open(f"{op.repo_path}/test.txt", "w") as f:
+        f.write("test")
+    op.git_cli.git.add(A=True)
+    op.git_cli.create_remote(name="other-remote", url=op.clone_url)
+
+    assert op.git_cli.active_branch.name == "test_branch"
+    assert len(op.git_cli.branches) == num_branches + 1
+    assert len(op.git_cli.remotes) == 2
+    assert op.git_cli.is_dirty()
+
+    op.clean_repo()
+    assert not op.git_cli.is_dirty()  # discards changes
+    assert len(op.git_cli.branches) == 1  # deletes only the checked out branch
+    assert op.git_cli.active_branch.name == op.default_branch
+    assert len(op.git_cli.remotes) == 1  # deletes all remotes except origin
+    assert op.git_cli.remotes[0].name == "origin"