Make codebase.reset only reset changes made by the sdk (#74)

# Motivation
<!-- Why is this change necessary? -->

# Content
<!-- Please include a summary of the change -->
# Testing
<!-- How was the change tested? -->
# Please check the following before marking your PR as ready for review

- [ ] I have added tests for my changes
- [ ] I have updated the documentation or added new documentation as
needed
- [ ] I have read and agree to the [Contributor License
Agreement](../CLA.md)

---------

Co-authored-by: bagel897 <ellenagarwal897@gmail.com>
Co-authored-by: codegen-bot <team+codegenbot@codegen.sh>
Co-authored-by: Carol Jung <165736129+caroljung-cg@users.noreply.github.com>
Co-authored-by: Rushil Patel <rpatel@codegen.com>
Co-authored-by: Christine Wang <christine@codegen.com>
Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com>
Co-authored-by: Jay Hack <jayhack@users.noreply.github.com>
Co-authored-by: tomcodgen <tkucar@codegen.com>
Co-authored-by: jemeza-codegen <165736868+jemeza-codegen@users.noreply.github.com>

Author: 10 people

.gitignore

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

pyproject.toml

@@ -139,6 +139,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://aws@codegen-922078275900.d.codeartifact.us-east-1.amazonaws.com/pypi/codegen/simple/"]

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):