fix: improve unknown branch detection (#18)

Author: mtth

src/git_draft/assistants/common.py

@@ -11,8 +11,15 @@ class Toolbox(Protocol):
     # signature-only versions.
 
     def list_files(self) -> Sequence[PurePosixPath]: ...
+
     def read_file(self, path: PurePosixPath) -> str: ...
-    def write_file(self, path: PurePosixPath, contents: str) -> None: ...
+
+    def write_file(
+        self,
+        path: PurePosixPath,
+        contents: str,
+        change_description: str | None = None,
+    ) -> None: ...
 
 
 @dataclasses.dataclass(frozen=True)

src/git_draft/assistants/openai.py

@@ -26,7 +26,7 @@ def _function_tool_param(
                 "type": "object",
                 "additionalProperties": False,
                 "properties": inputs or {},
-                "required": required_inputs or [],
+                "required": list(inputs.keys()) if inputs else [],
             },
             "strict": True,
         },
@@ -47,7 +47,6 @@ def _function_tool_param(
                 "description": "Path of the file to be read",
             },
         },
-        required_inputs=["path"],
     ),
     _function_tool_param(
         name="write_file",
@@ -65,8 +64,13 @@ def _function_tool_param(
                 "type": "string",
                 "description": "New contents of the file",
             },
+            "change_description": {
+                "type": "string",
+                "description": """\
+                    Brief description of the changes performed on this file
+                """,
+            },
         },
-        required_inputs=["path", "contents"],
     ),
 ]
 
@@ -77,6 +81,8 @@ def _function_tool_param(
     You are an expert software engineer, who writes correct and concise code.
     Use the provided functions to find the filesyou need to answer the query,
     read the content of the relevant ones, and save the changes you suggest.
+    When writing a file, include a summary description of the changes you have
+    made.
 """
 
 

src/git_draft/common.py

@@ -11,7 +11,7 @@
 import string
 import sys
 import tempfile
-from typing import ContextManager
+from typing import Iterator
 import xdg_base_dirs
 
 
@@ -30,9 +30,7 @@ def _guess_editor_binpath() -> str:
 
 
 def _get_tty_filename():
-    if sys.platform == "win32":
-        return "CON:"
-    return "/dev/tty"
+    return "CON:" if sys.platform == "win32" else "/dev/tty"
 
 
 def open_editor(placeholder="") -> str:
@@ -80,8 +78,16 @@ def persistent(cls) -> Store:
     def in_memory(cls) -> Store:
         return cls(sqlite3.connect(":memory:"))
 
-    def cursor(self) -> ContextManager[sqlite3.Cursor]:
-        return contextlib.closing(self._connection.cursor())
+    @contextlib.contextmanager
+    def cursor(self) -> Iterator[sqlite3.Cursor]:
+        with contextlib.closing(self._connection.cursor()) as cursor:
+            try:
+                yield cursor
+            except:  # noqa
+                self._connection.rollback()
+                raise
+            else:
+                self._connection.commit()
 
 
 _query_root = Path(__file__).parent / "queries"

src/git_draft/manager.py

@@ -63,7 +63,12 @@ def read_file(self, path: PurePosixPath) -> str:
         # Read the file from the index.
         return self._repo.git.show(f":{path}")
 
-    def write_file(self, path: PurePosixPath, contents: str) -> None:
+    def write_file(
+        self,
+        path: PurePosixPath,
+        contents: str,
+        change_description: str | None = None,
+    ) -> None:
         # Update the index without touching the worktree.
         # https://stackoverflow.com/a/25352119
         with tempfile.NamedTemporaryFile(delete_on_close=False) as temp:
@@ -187,9 +192,12 @@ def _exit_draft(self, apply: bool, delete=False) -> None:
             raise RuntimeError("Not currently on a draft branch")
 
         with self._store.cursor() as cursor:
-            [(origin_branch, origin_sha, sync_sha)] = cursor.execute(
+            rows = cursor.execute(
                 sql("get-branch-by-suffix"), {"suffix": branch.suffix}
             )
+            if not rows:
+                raise RuntimeError("Unrecognized branch")
+            [(origin_branch, origin_sha, sync_sha)] = rows
 
         if (
             not apply

tests/git_draft/common_test.py

@@ -17,6 +17,17 @@ def test_ensure_state_home(state_home) -> None:
     assert path.exists()
 
 
+class TestRandomId:
+    def test_length(self) -> None:
+        length = 10
+        result = sut.random_id(length)
+        assert len(result) == length
+
+    def test_content(self) -> None:
+        result = sut.random_id(1000)
+        assert set(result).issubset(sut._alphabet)
+
+
 class TestStore:
     def test_cursor(self) -> None:
         store = sut.Store.persistent()
@@ -30,3 +41,8 @@ def test_cursor(self) -> None:
 
 def test_sql() -> None:
     assert "create" in sut.sql("create-tables")
+
+
+def test_sql_missing() -> None:
+    with pytest.raises(FileNotFoundError):
+        sut.sql("non_existent_file")