Redo undo printing

mentat/code_feature.py

@@ -152,8 +152,8 @@ def __repr__(self):
         )
 
     def ref(self, cwd: Optional[Path] = None) -> str:
-        if cwd is not None and self.path.is_relative_to(cwd):
-            path_string = self.path.relative_to(cwd)
+        if cwd is not None:
+            path_string = str(get_relative_path(self.path, cwd))
         else:
             path_string = str(self.path)
 

mentat/code_file_manager.py

@@ -5,22 +5,15 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from mentat.edit_history import (
-    CreationAction,
-    DeletionAction,
-    EditAction,
-    EditHistory,
-    RenameAction,
-)
+from mentat.edit_history import EditHistory
 from mentat.errors import MentatError
 from mentat.interval import Interval
 from mentat.session_context import SESSION_CONTEXT
 from mentat.session_input import ask_yes_no
-from mentat.utils import sha256
+from mentat.utils import get_relative_path, sha256
 
 if TYPE_CHECKING:
     # This normally will cause a circular import
-    from mentat.code_context import CodeContext
     from mentat.parsers.file_edit import FileEdit
 
 
@@ -38,22 +31,29 @@ def read_file(self, path: Path) -> list[str]:
         self.file_lines[abs_path] = lines
         return lines
 
-    def _create_file(self, code_context: CodeContext, abs_path: Path):
+    def create_file(self, abs_path: Path, content: str = ""):
+        ctx = SESSION_CONTEXT.get()
+        code_context = ctx.code_context
+
         logging.info(f"Creating new file {abs_path}")
         # Create any missing directories in the path
         abs_path.parent.mkdir(parents=True, exist_ok=True)
         with open(abs_path, "w") as f:
-            f.write("")
+            f.write(content)
         code_context.include(abs_path)
 
-    def _delete_file(self, code_context: CodeContext, abs_path: Path):
+    def delete_file(self, abs_path: Path):
+        ctx = SESSION_CONTEXT.get()
+        code_context = ctx.code_context
+
         logging.info(f"Deleting file {abs_path}")
         code_context.exclude(abs_path)
         abs_path.unlink()
 
-    def _rename_file(
-        self, code_context: CodeContext, abs_path: Path, new_abs_path: Path
-    ):
+    def rename_file(self, abs_path: Path, new_abs_path: Path):
+        ctx = SESSION_CONTEXT.get()
+        code_context = ctx.code_context
+
         logging.info(f"Renaming file {abs_path} to {new_abs_path}")
         code_context.exclude(abs_path)
         os.rename(abs_path, new_abs_path)
@@ -66,27 +66,22 @@ async def write_changes_to_files(
     ) -> list[FileEdit]:
         session_context = SESSION_CONTEXT.get()
         stream = session_context.stream
-        code_context = session_context.code_context
         agent_handler = session_context.agent_handler
 
         if not file_edits:
             return []
 
         applied_edits: list[FileEdit] = []
         for file_edit in file_edits:
-            if file_edit.file_path.is_relative_to(session_context.cwd):
-                display_path = file_edit.file_path.relative_to(session_context.cwd)
-            else:
-                display_path = file_edit.file_path
+            display_path = get_relative_path(file_edit.file_path, session_context.cwd)
 
             if file_edit.is_creation:
                 if file_edit.file_path.exists():
                     raise MentatError(
                         f"Model attempted to create file {file_edit.file_path} which"
                         " already exists"
                     )
-                self.history.add_action(CreationAction(file_edit.file_path))
-                self._create_file(code_context, file_edit.file_path)
+                self.create_file(file_edit.file_path)
             elif not file_edit.file_path.exists():
                 raise MentatError(
                     f"Attempted to edit non-existent file {file_edit.file_path}"
@@ -100,12 +95,8 @@ async def write_changes_to_files(
                 if await ask_yes_no(default_yes=False):
                     stream.send(f"Deleting {display_path}...", color="red")
                     # We use the current lines rather than the stored lines for undo
-                    self.history.add_action(
-                        DeletionAction(
-                            file_edit.file_path, self.read_file(file_edit.file_path)
-                        )
-                    )
-                    self._delete_file(code_context, file_edit.file_path)
+                    file_edit.previous_file_lines = self.read_file(file_edit.file_path)
+                    self.delete_file(file_edit.file_path)
                     applied_edits.append(file_edit)
                     continue
                 else:
@@ -134,23 +125,19 @@ async def write_changes_to_files(
                         f"Attempted to rename file {file_edit.file_path} to existing"
                         f" file {file_edit.rename_file_path}"
                     )
-                self.history.add_action(
-                    RenameAction(file_edit.file_path, file_edit.rename_file_path)
-                )
-                self._rename_file(
-                    code_context, file_edit.file_path, file_edit.rename_file_path
-                )
-                file_edit.file_path = file_edit.rename_file_path
+                self.rename_file(file_edit.file_path, file_edit.rename_file_path)
 
             new_lines = file_edit.get_updated_file_lines(stored_lines)
             if new_lines != stored_lines:
+                file_path = file_edit.rename_file_path or file_edit.file_path
                 # We use the current lines rather than the stored lines for undo
-                self.history.add_action(
-                    EditAction(file_edit.file_path, self.read_file(file_edit.file_path))
-                )
-                with open(file_edit.file_path, "w") as f:
+                file_edit.previous_file_lines = self.read_file(file_path)
+                with open(file_path, "w") as f:
                     f.write("\n".join(new_lines))
             applied_edits.append(file_edit)
+
+        for applied_edit in applied_edits:
+            self.history.add_edit(applied_edit)
         if not agent_handler.agent_enabled:
             self.history.push_edits()
         return applied_edits

mentat/edit_history.py

@@ -1,107 +1,26 @@
-import os
-from pathlib import Path
 from typing import Optional
 
-import attr
 from termcolor import colored
 
 from mentat.errors import HistoryError
-from mentat.parsers.file_edit import FileEdit, Replacement
+from mentat.parsers.file_edit import FileEdit
 from mentat.session_context import SESSION_CONTEXT
 
 
-# All paths should be abs paths
-@attr.define()
-class RenameAction:
-    old_file_name: Path = attr.field()
-    cur_file_name: Path = attr.field()
-
-    def undo(self) -> FileEdit:
-        if self.old_file_name.exists():
-            raise HistoryError(
-                f"File {self.old_file_name} already exists; unable to undo rename from"
-                f" {self.cur_file_name}"
-            )
-        else:
-            os.rename(self.cur_file_name, self.old_file_name)
-            return FileEdit(
-                file_path=self.old_file_name, rename_file_path=self.cur_file_name
-            )
-
-
-@attr.define()
-class CreationAction:
-    cur_file_name: Path = attr.field()
-
-    def undo(self) -> FileEdit:
-        if not self.cur_file_name.exists():
-            raise HistoryError(
-                f"File {self.cur_file_name} does not exist; unable to delete"
-            )
-        else:
-            self.cur_file_name.unlink()
-            return FileEdit(file_path=self.cur_file_name, is_creation=True)
-
-
-@attr.define()
-class DeletionAction:
-    old_file_name: Path = attr.field()
-    old_file_lines: list[str] = attr.field()
-
-    def undo(self) -> FileEdit:
-        if self.old_file_name.exists():
-            raise HistoryError(
-                f"File {self.old_file_name} already exists; unable to re-create"
-            )
-        else:
-            with open(self.old_file_name, "w") as f:
-                f.write("\n".join(self.old_file_lines))
-            return FileEdit(file_path=self.old_file_name, is_deletion=True)
-
-
-@attr.define()
-class EditAction:
-    cur_file_name: Path = attr.field()
-    old_file_lines: list[str] = attr.field()
-
-    def undo(self) -> FileEdit:
-        if not self.cur_file_name.exists():
-            raise HistoryError(
-                f"File {self.cur_file_name} does not exist; unable to undo edit"
-            )
-        else:
-            new_file_lines = self.cur_file_name.read_text().split("\n")
-            with open(self.cur_file_name, "w") as f:
-                f.write("\n".join(self.old_file_lines))
-            return FileEdit(
-                file_path=self.cur_file_name,
-                replacements=[
-                    Replacement(
-                        starting_line=0,
-                        ending_line=len(self.old_file_lines),
-                        new_lines=new_file_lines,
-                    )
-                ],
-            )
-
-
-HistoryAction = RenameAction | CreationAction | DeletionAction | EditAction
-
-
 # TODO: Keep track of when we create directories so we can undo those as well
 class EditHistory:
     def __init__(self):
-        self.edits = list[list[HistoryAction]]()
-        self.cur_edit = list[HistoryAction]()
+        self.edits = list[list[FileEdit]]()
+        self.cur_edits = list[FileEdit]()
         self.undone_edits = list[list[FileEdit]]()
 
-    def add_action(self, history_action: HistoryAction):
-        self.cur_edit.append(history_action)
+    def add_edit(self, file_edit: FileEdit):
+        self.cur_edits.append(file_edit)
 
     def push_edits(self):
-        if self.cur_edit:
-            self.edits.append(self.cur_edit)
-            self.cur_edit = list[HistoryAction]()
+        if self.cur_edits:
+            self.edits.append(self.cur_edits)
+            self.cur_edits = list[FileEdit]()
 
     def undo(self) -> str:
         if not self.edits:
@@ -112,10 +31,10 @@ def undo(self) -> str:
         errors = list[str]()
         undone_edit = list[FileEdit]()
         while cur_edit:
-            cur_action = cur_edit.pop()
+            cur_file_edit = cur_edit.pop()
             try:
-                redo_edit = cur_action.undo()
-                undone_edit.append(redo_edit)
+                cur_file_edit.undo()
+                undone_edit.append(cur_file_edit)
             except HistoryError as e:
                 errors.append(colored(str(e), color="light_red"))
         if undone_edit:
@@ -131,6 +50,8 @@ async def redo(self) -> Optional[str]:
 
         edits_to_redo = self.undone_edits.pop()
         edits_to_redo.reverse()
+        for edit in edits_to_redo:
+            edit.display_full_edit(code_file_manager.file_lines[edit.file_path])
         await code_file_manager.write_changes_to_files(edits_to_redo)
 
     def undo_all(self) -> str:

mentat/parsers/change_display_helper.py

@@ -9,6 +9,9 @@
 from pygments.util import ClassNotFound
 from termcolor import colored
 
+from mentat.session_context import SESSION_CONTEXT
+from mentat.utils import get_relative_path
+
 change_delimiter = 60 * "="
 
 
@@ -58,8 +61,12 @@ class DisplayInformation:
     new_name: Path | None = attr.field(default=None)
 
     def __attrs_post_init__(self):
+        ctx = SESSION_CONTEXT.get()
+
         self.line_number_buffer = get_line_number_buffer(self.file_lines)
         self.lexer = _get_lexer(self.file_name)
+        if self.file_name.is_absolute():
+            self.file_name = get_relative_path(self.file_name, ctx.cwd)
 
 
 def _remove_extra_empty_lines(lines: list[str]) -> list[str]: