feat: implement accept option

Author: mtth

README.md

@@ -23,8 +23,8 @@ pipx install git-draft[openai]
 * Mechanism for reporting feedback from a bot, and possibly allowing user to
   interactively respond.
 * Add configuration option to auto sync and `--no-sync` flag. Similar to reset.
-* Add "amend" commit when finalizing. This could be useful training data,
-  showing what the bot did not get right.
-* Convenience `--accept` functionality for simple cases: checkout option which
-  applies the changes, and finalizes the draft if specified multiple times. For
-  example `git draft -aa add-test symbol=foo`
+* Add optional "amend" commit when finalizing. This could be useful training
+  data, showing what the bot did not get right.
+* Support file rename tool.
+  * https://stackoverflow.com/q/49853177/1062617
+  * https://stackoverflow.com/q/6658313/1062617

docs/git-draft.adoc

@@ -18,7 +18,9 @@ IMPORTANT: `git-draft` is WIP.
 == Synopsis
 
 [verse]
-git draft [options] [--generate] [--bot BOT] [--edit] [--reset | --no-reset] [--sync] [TEMPLATE [VARIABLE...]]
+git draft [options] [--generate] [--accept...] [--bot BOT] [--edit]
+                                 [--reset | --no-reset] [--sync]
+                                 [TEMPLATE [VARIABLE...]]
 git draft [options] --finalize [--clean | --revert] [--delete]
 git draft [options] --show-drafts [--json]
 git draft [options] --show-prompts [--json] [PROMPT]
@@ -27,34 +29,43 @@ git draft [options] --show-templates [--json | [--edit] TEMPLATE]
 
 == Description
 
-`git-draft` is a git-centric way to edit code using AI.
+`git-draft` is a git-centric way to develop using AI.
 
 
 == Options
 
+-a::
+--accept*::
+  Check out generated changes automatically.
+  Can be repeated.
+  This may fail if you manually edit files that the bot updates during generation.
+
+
 -b BOT::
 --bot=BOT::
   Bot name.
 
 -c::
---clean::
-  TODO
+--clean*::
+  Remove files deleted during a draft which otherwise would show up as untracked.
 
 -d::
 --delete::
   Delete finalized branch.
 
 -e::
 --edit::
-  Edit.
+  Enable interactive editing of prompts and templates.
+  See `--generate` and `--show-templates` for details.
 
 -F::
 --finalize::
-  TODO
+  Go back to the draft's origin branch with the current working directory.
 
 -G::
 --generate::
-  TODO
+  Add an AI-generated commit.
+  If the `--edit` option is set, an interactive editor will be open with the rendered prompt to allow modification before it is shared with the bot.
 
 -h::
 --help::
@@ -69,26 +80,29 @@ git draft [options] --show-templates [--json | [--edit] TEMPLATE]
 
 --reset::
 --no-reset::
-  TODO
+  Controls behavior when staged changes are present at the start of a generate command.
+   If enabled, these changes are automatically reset and combined with other working directory changes.
+   Otherwise an error is raised.
 
 -r::
---revert::
+--revert*::
   Abandon any changes since draft creation.
 
 --root::
   Repository search root.
 
 -D::
 --show-drafts::
-  TODO
+  List recently created drafts.
 
 -P::
 --show-prompts::
-  TODO
+  Lists recently used prompts.
 
 -T::
 --show-templates::
-  TODO
+  Lists available templates.
+  With an template name argument, displays the corresponding template's contents or, if the `--edit` option is set, opens an interactive editor.
 
 -s::
 --sync::
@@ -102,6 +116,10 @@ git draft [options] --show-templates [--json | [--edit] TEMPLATE]
   Show version and exit.
 
 
+Options suffixed with * may modify the working directory's contents.
+All other combinations of options never write to it, and are safe to do concurrently with other working directory operations.
+
+
 == Examples
 
 The workhorse command is `git draft --generate` which leverages AI to edit our code.

src/git_draft/__main__.py

@@ -11,7 +11,7 @@
 
 from .bots import load_bot
 from .common import PROGRAM, Config, UnreachableError, ensure_state_home
-from .drafter import Drafter
+from .drafter import Accept, Drafter
 from .editor import open_editor
 from .prompt import Template, TemplatedPrompt, templates_table
 from .store import Store
@@ -58,6 +58,12 @@ def callback(_option, _opt, _value, parser) -> None:
     add_command("show-prompts", short="P", help="show prompt history")
     add_command("show-templates", short="T", help="show template information")
 
+    parser.add_option(
+        "-a",
+        "--accept",
+        help="apply generated changes",
+        action="count",
+    )
     parser.add_option(
         "-b",
         "--bot",
@@ -204,6 +210,7 @@ def main() -> None:  # noqa: PLR0912 PLR0915
         name = drafter.generate_draft(
             prompt,
             bot,
+            accept=Accept(opts.accept or 0),
             bot_name=opts.bot,
             prompt_transform=open_editor if editable else None,
             tool_visitors=[ToolPrinter()],

src/git_draft/drafter.py

@@ -4,6 +4,7 @@
 
 import dataclasses
 from datetime import datetime
+import enum
 import json
 import logging
 import os
@@ -57,6 +58,14 @@ def new_suffix():
         return random_id(9)
 
 
+class Accept(enum.Enum):
+    """Valid accept modes"""
+
+    MANUAL = enum.auto()
+    CHECKOUT = enum.auto()
+    FINALIZE = enum.auto()
+
+
 class Drafter:
     """Draft state orchestrator"""
 
@@ -77,12 +86,13 @@ def generate_draft(  # noqa: PLR0913
         self,
         prompt: str | TemplatedPrompt,
         bot: Bot,
+        accept: Accept = Accept.MANUAL,
         bot_name: str | None = None,
-        tool_visitors: Sequence[ToolVisitor] | None = None,
         prompt_transform: Callable[[str], str] | None = None,
         reset: bool = False,
         sync: bool = False,
         timeout: float | None = None,
+        tool_visitors: Sequence[ToolVisitor] | None = None,
     ) -> str:
         if timeout is not None:
             raise NotImplementedError()  # TODO
@@ -172,7 +182,17 @@ def generate_draft(  # noqa: PLR0913
             )
 
         _logger.info("Completed generation for %s.", branch)
-        return str(branch)
+        if accept == Accept.MANUAL:
+            return str(branch)
+
+        # Check out files from the index. Since we assume that users do not
+        # manually update the index in draft branches, this is equivalent to
+        # checking out the files from the latest (generated, here) commit.
+        # delta = self._delta(
+        self._repo.git.checkout(".", theirs=True)
+        if accept == Accept.CHECKOUT:
+            return str(branch)
+        return self.exit_draft(revert=False, clean=accept == Accept.CLEAN)
 
     def exit_draft(self, *, revert: bool, clean=False, delete=False) -> str:
         branch = _Branch.active(self._repo)
@@ -195,9 +215,10 @@ def exit_draft(self, *, revert: bool, clean=False, delete=False) -> str:
             raise RuntimeError("Parent branch has moved, please rebase first")
 
         if clean and not revert:
-            # We delete files which have been deleted in the draft manually,
+            _logger.debug("Cleaning up files.")
+            # We manually delete files which have been deleted in the draft,
             # otherwise they would still show up as untracked.
-            origin_delta = self._delta(f"{origin_branch}..{branch}")
+            origin_delta = self._delta(start=origin_branch, end=str(branch))
             deleted = self._untracked() & origin_delta.deleted
             for path in deleted:
                 os.remove(osp.join(self._repo.working_dir, path))
@@ -211,17 +232,18 @@ def exit_draft(self, *, revert: bool, clean=False, delete=False) -> str:
         self._repo.git.checkout(origin_branch)
 
         if revert:
+            _logger.debug("Reverting changes... [sync_sha=%s]", sync_sha)
             # We revert the relevant files if needed. If a sync commit had been
             # created, we simply revert to it. Otherwise we compute which files
             # have changed due to draft commits and revert only those.
             if sync_sha:
-                delta = self._delta(sync_sha)
-                if delta.changed:
-                    self._repo.git.checkout(sync_sha, "--", ".")
+                self._repo.git.checkout("-f", sync_sha)
                 _logger.info("Reverted to sync commit. [sha=%s]", sync_sha)
             else:
-                origin_delta = self._delta(f"{origin_branch}..{branch}")
-                head_delta = self._delta("HEAD")
+                origin_delta = self._delta(
+                    start=origin_branch, end=str(branch)
+                )
+                head_delta = self._delta(end="HEAD")
                 changed = head_delta.touched & origin_delta.changed
                 if changed:
                     self._repo.git.checkout("--", *changed)
@@ -304,15 +326,18 @@ def _create_branch(self, sync: bool) -> _Branch:
     def _stage_changes(self, sync: bool) -> str | None:
         self._repo.git.add(all=True)
         if not sync or not self._repo.is_dirty(untracked_files=True):
+            _logger.debug("Skipped sync commit creation. [sync=%s]", sync)
             return None
         ref = self._repo.index.commit("draft! sync")
+        _logger.debug("Created sync commit. [sha=%s]", ref.hexsha)
         return ref.hexsha
 
     def _untracked(self) -> frozenset[str]:
         text = self._repo.git.ls_files(exclude_standard=True, others=True)
         return frozenset(text.splitlines())
 
-    def _delta(self, spec) -> _Delta:
+    def _delta(self, *, start: str | None = None, end: str) -> _Delta:
+        spec = f"{start}..{end}" if start else end
         changed = list[str]()
         deleted = list[str]()
         for line in self._repo.git.diff(spec, name_status=True).splitlines():
@@ -321,7 +346,9 @@ def _delta(self, spec) -> _Delta:
                 deleted.append(name)
             else:
                 changed.append(name)
-        return _Delta(changed=frozenset(changed), deleted=frozenset(deleted))
+        delta = _Delta(changed=frozenset(changed), deleted=frozenset(deleted))
+        _logger.debug("Computed delta for %s: %s", spec, delta)
+        return delta
 
 
 @dataclasses.dataclass(frozen=True)

tests/git_draft/drafter_test.py

@@ -100,7 +100,7 @@ def test_generate_dirty_index_reset_sync(self) -> None:
         assert self._read("PROMPT") == "hi"
         assert len(self._commits()) == 3  # init, sync, prompt
 
-    def test_generate_clean_index_sync(self) -> None:
+    def test_generate_empty_index_sync(self) -> None:
         prompt = TemplatedPrompt("add-test", {"symbol": "abc"})
         self._drafter.generate_draft(prompt, FakeBot(), sync=True)
         self._repo.git.checkout(".")
@@ -142,20 +142,45 @@ def act(self, _goal: Goal, toolbox: Toolbox) -> Action:
 
     def test_sync_delete_revert(self) -> None:
         self._write("p1", "a")
+        self._write("p2", "b")
         self._repo.git.add(all=True)
         self._repo.index.commit("advance")
         self._delete("p1")
+        self._delete("p2")
+        self._write("p3", "c")
 
         class CustomBot(Bot):
             def act(self, _goal: Goal, toolbox: Toolbox) -> Action:
                 toolbox.write_file(PurePosixPath("p2"), "b")
+                toolbox.write_file(PurePosixPath("p4"), "d")
                 return Action()
 
-        self._drafter.generate_draft("hello", CustomBot(), sync=True)
+        self._drafter.generate_draft(
+            "hello", CustomBot(), accept=sut.Accept.CHECKOUT, sync=True
+        )
         assert self._read("p1") is None
+        assert self._read("p2") == "b"
+        assert self._read("p3") == "c"
+        assert self._read("p4") == "d"
 
         self._drafter.exit_draft(revert=True)
         assert self._read("p1") is None
+        assert self._read("p2") is None
+        assert self._read("p3") == "c"
+        assert self._read("p4") is None
+
+    def test_generate_accept_deletion(self) -> None:
+        self._write("p1", "a")
+
+        class CustomBot(Bot):
+            def act(self, _goal: Goal, toolbox: Toolbox) -> Action:
+                toolbox.delete_file(PurePosixPath("p1"))
+                return Action()
+
+        self._drafter.generate_draft(
+            "hello", CustomBot(), accept=sut.Accept.CHECKOUT
+        )
+        assert self._read("p1") is None
 
     def test_generate_delete_finalize_clean(self) -> None:
         self._write("p1", "a")
@@ -167,10 +192,12 @@ def act(self, _goal: Goal, toolbox: Toolbox) -> Action:
                 toolbox.delete_file(PurePosixPath("p1"))
                 return Action()
 
-        self._drafter.generate_draft("hello", CustomBot())
-        assert self._read("p1") == "a"
+        self._drafter.generate_draft(
+            "hello", CustomBot(), accept=sut.Accept.CHECKOUT
+        )
+        assert self._read("p1") is None
 
-        self._drafter.exit_draft(revert=False, clean=True)
+        self._drafter.exit_draft(revert=False)
         assert self._read("p1") is None
 
     def test_revert_outside_draft(self) -> None: