feat: implement basic draft creation (#10)

Author: mtth

docs/git-draft.adoc

@@ -15,14 +15,12 @@ git-draft - git-friendly code assistant
 
 == Synopsis
 
-*git-draft* _-C_ _NAME_
+*git-draft* _-C_
 
-*git-draft* _-P_
+*git-draft* _-E_
 
 *git-draft* _-A_
 
-*git-draft* _-D_ _[-b BRANCH]_
-
 
 == Description
 
@@ -33,19 +31,15 @@ _git-draft_ is a git-centric way to edit code using AI.
 When you create a new draft with `git draft -C $name`, a new branch called `$branch/drafts/$name-$hash` is created (`$hash` is a random suffix used to guarantee uniqueness of branch names) and checked out.
 Additionally, any uncommitted changes are automatically committed (`draft! sync`).
 
-Once the draft is created, we can use AI to edit our code using `git draft -P`.
-It expects the prompt as standard input, for example `echo "Add a test for compute_offset in chart.py" | git draft -P`.
+Once the draft is created, we can use AI to edit our code using `git draft -E`.
+It expects the prompt as standard input, for example `echo "Add a test for compute_offset in chart.py" | git draft -E`.
 The prompt will automatically get augmented with information about the files in the repository, and give the AI access to tools for reading and writing files.
 Once the response has been received and changes, applied a commit is created (`draft! prompt: a short summary of the change`).
 
 The prompt step can be repeated as many times as needed. Once you are satisfied with the changes, run `git draft -A` to apply them.
 This will check out the branch used when creating the draft, adding the final state of the draft to the worktree.
 Note that you can come back to an existing draft anytime (by checking its branch out), but you will not be able to apply it if its origin branch has moved since the draft was created.
 
-Finally, calling `git draft -D` will delete all drafts associated with a branch.
-By default the currently active branch is used, you can use the `-b` option to select another.
-You can also delete the drafts manually by deleting their branches using `git branch`.
-
 
 == See also
 

poetry.lock

@@ -21,7 +21,9 @@ python = '>=3.12,<4'
 black = '^25.1.0'
 flake8 = '^7.0.0'
 flake8-pyproject = '^1.2.3'
+gitpython = '^3.1.44'
 mypy = '^1.2.0'
+openai = '^1.64.0'
 poethepoet = '^0.25.0'
 pytest = '^7.1.2'
 

pyproject.toml

@@ -1,3 +1,7 @@
-from .backend import Backend
+from .actions import apply_draft, create_draft, extend_draft
 
-__all__ = ["Backend"]
+__all__ = [
+    "apply_draft",
+    "create_draft",
+    "extend_draft",
+]

src/git_draft/__init__.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import optparse
+import sys
+
+from . import apply_draft, create_draft, extend_draft
 
 
 parser = optparse.OptionParser(prog="git-draft")
@@ -42,17 +45,14 @@ def option_group(self) -> optparse.OptionGroup:
             parser.add_option_group(self._option_group)
         return self._option_group
 
-    def __call__(self, _option, _opt, value, parser, name) -> None:
+    def __call__(self, _option, _opt, _value, parser, name) -> None:
         parser.values.command = name
-        parser.values.command_args = value
 
 
-Command.register(
-    "create", help="create a draft", type="string", metavar="NAME"
-)
+Command.register("create", help="create a draft")
 
 Command.register(
-    "prompt", help="read a prompt from stdin to add to the current draft"
+    "extend", help="read a prompt from stdin to add to the current draft"
 )
 
 apply_command = Command.register(
@@ -79,11 +79,12 @@ def main() -> None:
     (opts, args) = parser.parse_args()
     command = getattr(opts, "command", None)
     if command == "create":
-        print("Creating draft...")
-    elif command == "prompt":
-        print("Updating draft...")
+        create_draft()
+    elif command == "extend":
+        prompt = sys.stdin.read()
+        extend_draft(prompt)
     elif command == "apply":
-        print("Applying draft...")
+        apply_draft()
     elif command == "delete":
         print("Deleting draft...")
     else:

src/git_draft/__main__.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import dataclasses
+import git
+import random
+from pathlib import PurePosixPath
+import re
+import string
+import tempfile
+from typing import Match, Sequence
+
+from .backend import NewFileBackend
+
+
+def _enclosing_repo() -> git.Repo:
+    return git.Repo(search_parent_directories=True)
+
+
+_random = random.Random()
+
+_SUFFIX_LENGTH = 8
+
+_branch_name_pattern = re.compile(r"drafts/(.+)/(\w+)")
+
+
+@dataclasses.dataclass(frozen=True)
+class _DraftBranch:
+    parent: str
+    suffix: str
+    repo: git.Repo
+
+    def __str__(self) -> str:
+        return f"drafts/{self.parent}/{self.suffix}"
+
+    @classmethod
+    def create(cls, repo: git.Repo) -> _DraftBranch:
+        if not repo.active_branch:
+            raise RuntimeError("No currently active branch")
+        suffix = "".join(
+            _random.choice(string.ascii_lowercase + string.digits)
+            for _ in range(_SUFFIX_LENGTH)
+        )
+        return cls(repo.active_branch.name, suffix, repo)
+
+    @classmethod
+    def active(cls, repo: git.Repo) -> _DraftBranch:
+        match: Match | None = None
+        if repo.active_branch:
+            match = _branch_name_pattern.fullmatch(repo.active_branch.name)
+        if not match:
+            raise RuntimeError("Not currently on a draft branch")
+        return _DraftBranch(match[1], match[2], repo)
+
+
+@dataclasses.dataclass(frozen=True)
+class _CommitNotes:
+    # https://stackoverflow.com/a/40496777
+    pass
+
+
+def create_draft() -> None:
+    repo = _enclosing_repo()
+    draft_branch = _DraftBranch.create(repo)
+    ref = repo.create_head(str(draft_branch))
+    repo.git.checkout(ref)
+
+
+class _Toolbox:
+    def __init__(self, repo: git.Repo) -> None:
+        self._repo = repo
+
+    def list_files(self) -> Sequence[PurePosixPath]:
+        # Show staged files.
+        return self._repo.git.ls_files()
+
+    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, data: str) -> None:
+        # Update the index without touching the worktree.
+        # https://stackoverflow.com/a/25352119
+        with tempfile.NamedTemporaryFile(delete_on_close=False) as temp:
+            temp.write(data.encode("utf8"))
+            temp.close()
+            sha = self._repo.git.hash_object("-w", "--path", path, temp.name)
+            mode = 644  # TODO: Read from original file if it exists.
+            self._repo.git.update_index(
+                "--add", "--cacheinfo", f"{mode},{sha},{path}"
+            )
+
+
+def extend_draft(prompt: str) -> None:
+    repo = _enclosing_repo()
+    _ = _DraftBranch.active(repo)
+
+    if repo.is_dirty():
+        repo.git.add(all=True)
+        repo.index.commit("draft! sync")
+
+    NewFileBackend().run(_Toolbox(repo))
+
+
+def apply_draft(delete=False) -> None:
+    repo = _enclosing_repo()
+    branch = _DraftBranch.active(repo)
+
+    # TODO: Check that parent has not moved. We could do this for example by
+    # adding a note to the draft branch with the original branch's commit ref.
+
+    # https://stackoverflow.com/a/15993574
+    repo.git.checkout("--detach")
+    repo.git.reset("--soft", branch.parent)
+    repo.git.checkout(branch.parent)
+
+    if delete:
+        repo.git.branch("-D", str(branch))

src/git_draft/actions.py

@@ -1,2 +1,26 @@
+from pathlib import PurePosixPath
+from typing import Protocol, Sequence
+
+
+class Toolbox(Protocol):
+    def list_files(self) -> Sequence[PurePosixPath]: ...
+    def read_file(self, path: PurePosixPath) -> str: ...
+    def write_file(self, path: PurePosixPath, data: str) -> None: ...
+
+
 class Backend:
-    pass
+    def run(self, toolbox: Toolbox) -> None: ...
+
+
+class NewFileBackend(Backend):
+    def run(self, toolbox: Toolbox) -> None:
+        # send request to backend...
+        import time
+
+        time.sleep(2)
+
+        # Add files to index.
+        import random
+
+        name = f"foo-{random.randint(1, 100)}"
+        toolbox.write_file(PurePosixPath(name), "hello!\n")