Make codebase.reset only reset changes made by the sdk

.gitignore

@@ -65,3 +65,4 @@ graph-sitter-types/typings/**
 coverage.json
 tests/integration/verified_codemods/codemod_data/repo_commits.json
 .codegen/*
+.benchmarks/*

pyproject.toml

@@ -107,6 +107,8 @@ dev-dependencies = [
     "black>=24.8.0",
     "isort>=5.13.2",
     "emoji>=2.14.0",
+    "pytest-benchmark[histogram]>=5.1.0",
+    "loguru>=0.7.3",
 ]
 keyring-provider = "subprocess"
 #extra-index-url = ["https://[email protected]/pypi/codegen/simple/"]
@@ -197,7 +199,7 @@ pythonpath = "."
 norecursedirs = "repos expected"
 # addopts = -v --cov=app --cov-report=term
 
-addopts = "--dist=loadgroup --junitxml=build/test-results/test/TEST.xml --strict-config --import-mode=importlib --cov-context=test --cov-report=json --cov-config=pyproject.toml -p no:doctest"
+addopts = "--dist=loadgroup --junitxml=build/test-results/test/TEST.xml --strict-config --import-mode=importlib --cov-context=test --cov-report=json --cov-config=pyproject.toml -p no:doctest --benchmark-autosave"
 filterwarnings ="""
     ignore::DeprecationWarning:botocore.*:
     ignore::DeprecationWarning:sqlalchemy.*:

src/codegen/git/repo_operator/local_repo_operator.py

@@ -72,8 +72,8 @@ def create_from_files(cls, repo_path: str, files: dict[str, str], bot_commit: bo
     def create_from_commit(cls, repo_path: str, commit: str, url: str) -> Self:
         """Do a shallow checkout of a particular commit to get a repository from a given remote URL."""
         op = cls(repo_config=BaseRepoConfig(), repo_path=repo_path, bot_commit=False)
+        op.discard_changes()
         if op.get_active_branch_or_commit() != commit:
-            op.discard_changes()
             op.create_remote("origin", url)
             op.git_cli.remotes["origin"].fetch(commit, depth=1)
             op.checkout_commit(commit)

src/codegen/git/repo_operator/remote_repo_operator.py

@@ -42,8 +42,9 @@ def __init__(
         setup_option: SetupOption = SetupOption.PULL_OR_CLONE,
         shallow: bool = True,
         github_type: GithubType = GithubType.GithubEnterprise,
+        bot_commit: bool = True,
     ) -> None:
-        super().__init__(repo_config=repo_config, base_dir=base_dir)
+        super().__init__(repo_config=repo_config, base_dir=base_dir, bot_commit=bot_commit)
         self.github_type = github_type
         self.setup_repo_dir(setup_option=setup_option, shallow=shallow)
 

src/codegen/git/repo_operator/repo_operator.py

@@ -67,12 +67,21 @@ def viz_file_path(self) -> str:
     def git_cli(self) -> GitCLI:
         """Note: this is recursive, may want to look out"""
         git_cli = GitCLI(self.repo_path)
+        has_username = False
+        has_email = False
+        with git_cli.config_reader(None) as reader:
+            if reader.has_option("user", "name"):
+                has_username = True
+            if reader.has_option("user", "email"):
+                has_email = True
         with git_cli.config_writer("repository") as writer:
-            if self.bot_commit:
+            if not has_username or not has_email or self.bot_commit:
                 if not writer.has_section("user"):
                     writer.add_section("user")
-                writer.set("user", "name", CODEGEN_BOT_NAME)
-                writer.set("user", "email", CODEGEN_BOT_EMAIL)
+                if not has_username or self.bot_commit:
+                    writer.set("user", "name", CODEGEN_BOT_NAME)
+                if not has_email or self.bot_commit:
+                    writer.set("user", "email", CODEGEN_BOT_EMAIL)
         return git_cli
 
     @property

src/codegen/sdk/codebase/codebase_graph.py

@@ -30,6 +30,7 @@
 from codegen.sdk.core.interfaces.importable import Importable
 from codegen.sdk.core.node_id_factory import NodeId
 from codegen.sdk.enums import Edge, EdgeType, NodeType, ProgrammingLanguage
+from codegen.sdk.extensions.io import write_changes
 from codegen.sdk.extensions.sort import sort_editables
 from codegen.sdk.extensions.utils import uncache_all
 from codegen.sdk.typescript.external.ts_declassify.ts_declassify import TSDeclassify
@@ -107,6 +108,7 @@ class CodebaseGraph:
     flags: Flags
     session_options: SessionOptions = SessionOptions()
     projects: list[ProjectConfig]
+    unapplied_diffs: list[DiffLite]
 
     def __init__(
         self,
@@ -161,6 +163,7 @@ def __init__(
             self.synced_commit = None
         self.pending_syncs = []
         self.all_syncs = []
+        self.unapplied_diffs = []
         self.pending_files = set()
         self.flags = Flags()
 
@@ -232,9 +235,40 @@ def apply_diffs(self, diff_list: list[DiffLite]) -> None:
         self.generation += 1
         self._process_diff_files(by_sync_type)
 
+    def _reset_files(self, syncs: list[DiffLite]) -> None:
+        files_to_write = []
+        files_to_remove = []
+        modified_files = set()
+        for sync in syncs:
+            if sync.path in modified_files:
+                continue
+            if sync.change_type == ChangeType.Removed:
+                files_to_write.append((sync.path, sync.old_content))
+                modified_files.add(sync.path)
+                logger.info(f"Removing {sync.path} from disk")
+            elif sync.change_type == ChangeType.Modified:
+                files_to_write.append((sync.path, sync.old_content))
+                modified_files.add(sync.path)
+            elif sync.change_type == ChangeType.Renamed:
+                files_to_write.append((sync.rename_from, sync.old_content))
+                files_to_remove.append(sync.rename_to)
+                modified_files.add(sync.rename_from)
+                modified_files.add(sync.rename_to)
+            elif sync.change_type == ChangeType.Added:
+                files_to_remove.append(sync.path)
+                modified_files.add(sync.path)
+        logger.info(f"Writing {len(files_to_write)} files to disk and removing {len(files_to_remove)} files")
+        write_changes(files_to_remove, files_to_write)
+
+    @stopwatch
+    def reset_codebase(self) -> None:
+        self._reset_files(self.all_syncs + self.pending_syncs + self.unapplied_diffs)
+        self.unapplied_diffs.clear()
+
     @stopwatch
     def undo_applied_diffs(self) -> None:
         self.transaction_manager.clear_transactions()
+        self.reset_codebase()
         self.check_changes()
         self.pending_syncs.clear()  # Discard pending changes
         if len(self.all_syncs) > 0:
@@ -256,6 +290,9 @@ def _revert_diffs(self, diff_list: list[DiffLite]) -> None:
 
     def save_commit(self, commit: GitCommit) -> None:
         if commit is not None:
+            logger.info(f"Saving commit {commit.hexsha} to graph")
+            self.all_syncs.clear()
+            self.unapplied_diffs.clear()
             self.synced_commit = commit
             if self.config.feature_flags.verify_graph:
                 self.old_graph = self._graph.copy()
@@ -630,9 +667,11 @@ def commit_transactions(self, sync_graph: bool = True, sync_file: bool = True, f
         # Commit transactions for all contexts
         files_to_lock = self.transaction_manager.to_commit(files)
         diffs = self.transaction_manager.commit(files_to_lock)
-        # Filter diffs to only include files that are still in the graph
-        diffs = [diff for diff in diffs if self.get_file(diff.path) is not None]
-        self.pending_syncs.extend(diffs)
+        for diff in diffs:
+            if self.get_file(diff.path) is None:
+                self.unapplied_diffs.append(diff)
+            else:
+                self.pending_syncs.append(diff)
 
         # Write files if requested
         if sync_file:

src/codegen/sdk/codebase/diff_lite.py

@@ -1,4 +1,4 @@
-from enum import Enum, auto
+from enum import IntEnum, auto
 from os import PathLike
 from pathlib import Path
 from typing import NamedTuple, Self
@@ -7,7 +7,7 @@
 from watchfiles import Change
 
 
-class ChangeType(Enum):
+class ChangeType(IntEnum):
     Modified = auto()
     Removed = auto()
     Renamed = auto()
@@ -40,8 +40,9 @@ class DiffLite(NamedTuple):
 
     change_type: ChangeType
     path: Path
-    rename_from: str | None = None
-    rename_to: str | None = None
+    rename_from: Path | None = None
+    rename_to: Path | None = None
+    old_content: bytes | None = None
 
     @classmethod
     def from_watch_change(cls, change: Change, path: PathLike) -> Self:
@@ -52,11 +53,15 @@ def from_watch_change(cls, change: Change, path: PathLike) -> Self:
 
     @classmethod
     def from_git_diff(cls, git_diff: Diff):
+        old = None
+        if git_diff.a_blob:
+            old = git_diff.a_blob.data_stream.read()
         return cls(
             change_type=ChangeType.from_git_change_type(git_diff.change_type),
-            path=Path(git_diff.a_path),
-            rename_from=git_diff.rename_from,
-            rename_to=git_diff.rename_to,
+            path=Path(git_diff.a_path) if git_diff.a_path else None,
+            rename_from=Path(git_diff.rename_from) if git_diff.rename_from else None,
+            rename_to=Path(git_diff.rename_to) if git_diff.rename_to else None,
+            old_content=old,
         )
 
     @classmethod

src/codegen/sdk/codebase/transaction_manager.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from codegen.sdk.codebase.diff_lite import DiffLite
+from codegen.sdk.codebase.diff_lite import ChangeType, DiffLite
 from codegen.sdk.codebase.transactions import (
     EditTransaction,
     FileAddTransaction,
@@ -163,16 +163,16 @@ def to_commit(self, files: set[Path] | None = None) -> set[Path]:
             return set(self.queued_transactions.keys())
         return files.intersection(self.queued_transactions)
 
-    def commit(self, files: set[Path]) -> set[DiffLite]:
+    def commit(self, files: set[Path]) -> list[DiffLite]:
         """Execute transactions in bulk for each file, in reverse order of start_byte.
-        Returns the set of diffs that were committed.
+        Returns the list of diffs that were committed.
         """
         if self._commiting:
             logger.warn("Skipping commit, already committing")
-            return set()
+            return []
         self._commiting = True
         try:
-            diffs: set[DiffLite] = set()
+            diffs: list[DiffLite] = []
             if not self.queued_transactions or len(self.queued_transactions) == 0:
                 return diffs
 
@@ -187,9 +187,16 @@ def commit(self, files: set[Path]) -> set[DiffLite]:
                     logger.info(f"Committing {len(self.queued_transactions[file])} transactions for {file}")
             for file_path in files:
                 file_transactions = self.queued_transactions.pop(file_path, [])
+                modified = False
                 for transaction in file_transactions:
                     # Add diff IF the file is a source file
-                    diffs.add(transaction.get_diff())
+                    diff = transaction.get_diff()
+                    if diff.change_type == ChangeType.Modified:
+                        if not modified:
+                            modified = True
+                            diffs.append(diff)
+                    else:
+                        diffs.append(diff)
                     transaction.execute()
             return diffs
         finally:

src/codegen/sdk/codebase/transactions.py

@@ -127,7 +127,7 @@ def execute(self) -> None:
 
     def get_diff(self) -> DiffLite:
         """Gets the diff produced by this transaction"""
-        return DiffLite(ChangeType.Modified, self.file_path)
+        return DiffLite(ChangeType.Modified, self.file_path, old_content=self.file.content_bytes)
 
     def diff_str(self) -> str:
         """Human-readable string representation of the change"""
@@ -170,7 +170,7 @@ def execute(self) -> None:
 
     def get_diff(self) -> DiffLite:
         """Gets the diff produced by this transaction"""
-        return DiffLite(ChangeType.Modified, self.file_path)
+        return DiffLite(ChangeType.Modified, self.file_path, old_content=self.file.content_bytes)
 
     def diff_str(self) -> str:
         """Human-readable string representation of the change"""
@@ -205,7 +205,7 @@ def execute(self) -> None:
 
     def get_diff(self) -> DiffLite:
         """Gets the diff produced by this transaction"""
-        return DiffLite(ChangeType.Modified, self.file_path)
+        return DiffLite(ChangeType.Modified, self.file_path, old_content=self.file.content_bytes)
 
     def diff_str(self) -> str:
         """Human-readable string representation of the change"""
@@ -269,7 +269,7 @@ def execute(self) -> None:
 
     def get_diff(self) -> DiffLite:
         """Gets the diff produced by this transaction"""
-        return DiffLite(ChangeType.Renamed, self.file_path, self.file_path, self.new_file_path)
+        return DiffLite(ChangeType.Renamed, self.file_path, self.file_path, self.new_file_path, old_content=self.file.content_bytes)
 
     def diff_str(self) -> str:
         """Human-readable string representation of the change"""
@@ -294,7 +294,7 @@ def execute(self) -> None:
 
     def get_diff(self) -> DiffLite:
         """Gets the diff produced by this transaction"""
-        return DiffLite(ChangeType.Removed, self.file_path)
+        return DiffLite(ChangeType.Removed, self.file_path, old_content=self.file.content_bytes)
 
     def diff_str(self) -> str:
         """Human-readable string representation of the change"""

src/codegen/sdk/core/codebase.py

@@ -738,7 +738,7 @@ def current_commit(self) -> GitCommit | None:
         return self._op.git_cli.head.commit
 
     @stopwatch
-    def reset(self) -> None:
+    def reset(self, git_reset: bool = False) -> None:
         """Resets the codebase by:
         - Discarding any staged/unstaged changes
         - Resetting stop codemod limits: (max seconds, max transactions, max AI requests)
@@ -751,7 +751,8 @@ def reset(self) -> None:
         - .ipynb files (Jupyter notebooks, where you are likely developing)
         """
         logger.info("Resetting codebase ...")
-        self._op.discard_changes()  # Discard any changes made to the raw file state
+        if git_reset:
+            self._op.discard_changes()  # Discard any changes made to the raw file state
         self._num_ai_requests = 0
         self.reset_logs()
         self.G.undo_applied_diffs()
@@ -818,12 +819,14 @@ def get_diffs(self, base: str | None = None) -> list[Diff]:
         return self._op.get_diffs(base)
 
     @noapidoc
-    def get_diff(self, base: str | None = None) -> str:
+    def get_diff(self, base: str | None = None, stage_files: bool = False) -> str:
         """Produce a single git diff for all files."""
-        self._op.git_cli.git.add(A=True)  # add all changes to the index so untracked files are included in the diff
+        if stage_files:
+            self._op.git_cli.git.add(A=True)  # add all changes to the index so untracked files are included in the diff
         if base is None:
-            return self._op.git_cli.git.diff(patch=True, full_index=True, staged=True)
-        return self._op.git_cli.git.diff(base, full_index=True)
+            diff = self._op.git_cli.git.diff("HEAD", patch=True, full_index=True)
+            return diff
+        return self._op.git_cli.git.diff(base, patch=True, full_index=True)
 
     @noapidoc
     def clean_repo(self):

src/codegen/sdk/extensions/io.pyx

@@ -0,0 +1,12 @@
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
+import os
+
+
+def write_changes(files_to_remove: list[Path], files_to_write: list[tuple[Path, bytes]]):
+    # Start at the oldest sync and then apply non-conflicting newer changes
+    with ThreadPoolExecutor() as executor:
+        for file_to_remove in files_to_remove:
+            executor.submit(os.remove, file_to_remove)
+        for file_to_write, content in files_to_write:
+            executor.submit(file_to_write.write_bytes, content)

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

@@ -12,7 +12,8 @@
 
 @pytest.fixture
 def op(repo_config, request, tmpdir):
-    yield RemoteRepoOperator(repo_config, shallow=request.param, base_dir=tmpdir)
+    op = RemoteRepoOperator(repo_config, shallow=request.param, base_dir=tmpdir, bot_commit=False)
+    yield op
 
 
 @pytest.mark.parametrize("op", shallow_options, ids=lambda x: f"shallow={x}", indirect=True)

tests/integration/codemod/conftest.py

@@ -156,8 +156,9 @@ def _codebase(repo: Repo, op: RepoOperator, request) -> YieldFixture[Codebase]:
         projects = [ProjectConfig(repo_operator=op, programming_language=repo.language, subdirectories=repo.subdirectories, base_path=repo.base_path)]
         Codebases[repo.name] = Codebase(projects=projects, config=CodebaseConfig(feature_flags=feature_flags))
     codebase = Codebases[repo.name]
-    codebase.reset()
+    codebase.reset(git_reset=True)
     yield codebase
+    codebase.reset(git_reset=True)
 
 
 @pytest.fixture

tests/shared/codemod/codebase_comparison_utils.py

@@ -76,7 +76,7 @@ def compare_codebase_diff(
     diff = codebase.get_diff() + "\n"
     if not snapshot._snapshot_update:
         modified = gather_modified_files(codebase)
-        codebase.reset()
+        codebase.reset(git_reset=True)
         logger.info("Converting diff file to expected repository")
         if convert_diff_to_repo(expected_dir, expected_diff, codebase):
             return compare_codebase_with_snapshot(codebase, expected_dir, diff_path, snapshot, modified)