fixup! feat: add docstring prompt

Author: mtth

src/git_draft/__main__.py

@@ -9,13 +9,13 @@
 from .common import (
     Config,
     PROGRAM,
-    PromptRenderer,
     Store,
     UnreachableError,
     ensure_state_home,
     open_editor,
 )
 from .manager import Manager
+from .prompt import TemplatedPrompt
 
 
 def new_parser() -> optparse.OptionParser:
@@ -134,9 +134,7 @@ def main() -> None:
         prompt = opts.prompt
         if not prompt:
             if opts.template:
-                renderer = PromptRenderer.default()
-                kwargs = dict(e.split("=", 1) for e in args)
-                prompt = renderer.render(opts.template, **kwargs)
+                prompt = TemplatedPrompt.parse(opts.template, *args)
             elif sys.stdin.isatty():
                 prompt = open_editor("Enter your prompt here...")
             else:

src/git_draft/common.py

@@ -4,7 +4,6 @@
 import dataclasses
 from datetime import datetime
 import functools
-import jinja2
 import logging
 import os
 from pathlib import Path
@@ -27,7 +26,7 @@
 type JSONObject = Mapping[str, JSONValue]
 
 
-_package_root = Path(__file__).parent
+package_root = Path(__file__).parent
 
 
 def ensure_state_home() -> Path:
@@ -72,31 +71,6 @@ class BotConfig:
     pythonpath: str | None = None
 
 
-_prompt_root = _package_root / "prompts"
-
-
-class PromptRenderer:
-    def __init__(self, env: jinja2.Environment) -> None:
-        self._environment = env
-
-    @classmethod
-    def default(cls):
-        env = jinja2.Environment(
-            auto_reload=False,
-            autoescape=False,
-            keep_trailing_newline=True,
-            loader=jinja2.FileSystemLoader(
-                [Config.folder_path() / "prompts", str(_prompt_root)]
-            ),
-            undefined=jinja2.StrictUndefined,
-        )
-        return cls(env)
-
-    def render(self, template_name: str, **kwargs) -> str:
-        template = self._environment.get_template(f"{template_name}.jinja")
-        return template.render(kwargs)
-
-
 _default_editors = ["vim", "emacs", "nano"]
 
 
@@ -169,7 +143,7 @@ def cursor(self) -> Iterator[sqlite3.Cursor]:
                 self._connection.commit()
 
 
-_query_root = _package_root / "queries"
+_query_root = package_root / "queries"
 
 
 @functools.cache

src/git_draft/manager.py

@@ -13,6 +13,7 @@
 
 from .bots import Bot, OperationHook, Toolbox
 from .common import Store, random_id, sql
+from .prompt import PromptRenderer, TemplatedPrompt
 
 
 _logger = logging.getLogger(__name__)
@@ -135,13 +136,13 @@ def _create_branch(self, sync: bool) -> _Branch:
 
     def generate_draft(
         self,
-        prompt: str,
+        prompt: str | TemplatedPrompt,
         bot: Bot,
         checkout=False,
         reset=False,
         sync=False,
     ) -> None:
-        if not prompt.strip():
+        if isinstance(prompt, str) and not prompt.strip():
             raise ValueError("Empty prompt")
         if self._repo.is_dirty(working_tree=False):
             if not reset:
@@ -157,24 +158,31 @@ def generate_draft(
             branch = self._create_branch(sync)
             _logger.debug("Created branch %s.", branch)
 
+        if isinstance(prompt, TemplatedPrompt):
+            renderer = PromptRenderer.for_repo(self._repo)
+            prompt_contents = renderer.render(prompt)
+        else:
+            prompt_contents = prompt
         with self._store.cursor() as cursor:
             [(prompt_id,)] = cursor.execute(
                 sql("add-prompt"),
                 {
                     "branch_suffix": branch.suffix,
-                    "contents": prompt,
+                    "contents": prompt_contents,
                 },
             )
 
         start_time = time.perf_counter()
         toolbox = _Toolbox(self._repo, self._operation_hook)
-        action = bot.act(prompt, toolbox)
+        action = bot.act(prompt_contents, toolbox)
         end_time = time.perf_counter()
 
         title = action.title
         if not title:
-            title = textwrap.shorten(prompt, break_on_hyphens=False, width=72)
-        commit = self._repo.index.commit(f"draft! {title}\n\n{prompt}")
+            title = _default_title(prompt_contents)
+        commit = self._repo.index.commit(
+            f"draft! {title}\n\n{prompt_contents}"
+        )
 
         with self._store.cursor() as cursor:
             cursor.execute(
@@ -247,3 +255,7 @@ def _exit_draft(self, apply: bool, delete=False) -> None:
             self._repo.git.checkout(sync_sha, "--", ".")
         if delete:
             self._repo.git.branch("-D", branch.name)
+
+
+def _default_title(prompt: str) -> str:
+    return textwrap.shorten(prompt, break_on_hyphens=False, width=72)

src/git_draft/prompt.py

@@ -0,0 +1,44 @@
+import dataclasses
+import git
+import jinja2
+from typing import Mapping, Self
+
+from .common import Config, package_root
+
+
+_prompt_root = package_root / "prompts"
+
+
+@dataclasses.dataclass(frozen=True)
+class TemplatedPrompt:
+    template: str
+    context: Mapping[str, str]
+
+    @classmethod
+    def parse(cls, name: str, *args: str) -> Self:
+        return cls(name, dict(e.split("=", 1) for e in args))
+
+
+class PromptRenderer:
+    def __init__(self, env: jinja2.Environment) -> None:
+        self._environment = env
+
+    @classmethod
+    def for_repo(cls, repo: git.Repo):
+        env = jinja2.Environment(
+            auto_reload=False,
+            autoescape=False,
+            keep_trailing_newline=True,
+            loader=jinja2.FileSystemLoader(
+                [Config.folder_path() / "prompts", str(_prompt_root)]
+            ),
+            undefined=jinja2.StrictUndefined,
+        )
+        env.globals["repo"] = {
+            "file_paths": repo.git.ls_files(),
+        }
+        return cls(env)
+
+    def render(self, prompt: TemplatedPrompt) -> str:
+        template = self._environment.get_template(f"{prompt.template}.jinja")
+        return template.render(prompt.context)

src/git_draft/prompts/add-docstrings.jinja

@@ -1,5 +1,10 @@
-Add docstrings to all public functions and classes in {{ src_path }}. Be
-concise and do not repeat yourself.
+{% if src_path is defined %}
+Add docstrings to all public functions and classes in {{ src_path }}.
+{% else %}
+Add docstrings to all public functions and classes in this repository.
+{% endif %}
+
+Be concise and do not repeat yourself.
 
 Focus on highlighting aspects which are not obvious from the name of the
 symbols. Take time to look at the implementation to discover any behaviors
@@ -11,16 +16,18 @@ examples:
 
 ```python
 def write_file(path: Path, contents: str) -> None:
-  """Update a file's contents
+    """Update a file's contents
 
-  Args:
-    path: Path to the file to update.
-    contents: New file contents.
-  """
+    Args:
+        path: Path to the file to update.
+        contents: New file contents.
+    """
   ...
 
 class Renderer:
     """A smart renderer"""
 
     ...
 ```
+
+{% include "includes/file-list.jinja" %}

src/git_draft/prompts/add-test.jinja

@@ -8,9 +8,4 @@ tests. For example, if the surrounding code uses fixtures, do so as well.
 Do not stop until you have added at least one test. You should add separate
 tests to cover the normal execution path, and to cover any exceptional cases.
 
-For reference, here is the list of all currently available files in the
-repository:
-
-{% for path in repo.file_paths %}
-  * {{ path }}
-{% endfor %}
+{% include "includes/file-list.jinja" %}

src/git_draft/prompts/includes/file-list.jinja

@@ -0,0 +1,6 @@
+For reference, here is the list of all currently available files in the
+repository:
+
+{% for path in repo.file_paths %}
+  * {{ path }}
+{% endfor %}

tests/git_draft/common_test.py

@@ -85,13 +85,3 @@ def test_load_ok(self) -> None:
     def test_load_default(self) -> None:
         config = sut.Config.load()
         assert config.log_level == logging.INFO
-
-
-class TestPromptRenderer:
-    @pytest.fixture(autouse=True)
-    def setup(self) -> None:
-        self._renderer = sut.PromptRenderer.default()
-
-    def test_ok(self) -> None:
-        prompt = self._renderer.render("add-test", symbol="foo")
-        assert "foo" in prompt

tests/git_draft/conftest.py

@@ -0,0 +1,11 @@
+from pathlib import Path
+from typing import Iterator
+import git
+import pytest
+
+
+@pytest.fixture
+def repo(tmp_path: Path) -> Iterator[git.Repo]:
+    repo = git.Repo.init(str(tmp_path / "repo"), initial_branch="main")
+    repo.index.commit("init")
+    yield repo

tests/git_draft/manager_test.py

@@ -2,22 +2,12 @@
 import os.path as osp
 from pathlib import PurePosixPath
 import pytest
-import tempfile
-from typing import Iterator
 
 from git_draft.bots import Action, Bot, Toolbox
 from git_draft.common import Store
 import git_draft.manager as sut
 
 
-@pytest.fixture
-def repo() -> Iterator[git.Repo]:
-    with tempfile.TemporaryDirectory() as name:
-        repo = git.Repo.init(name, initial_branch="main")
-        repo.index.commit("init")
-        yield repo
-
-
 class _FakeBot(Bot):
     def act(self, prompt: str, toolbox: Toolbox) -> Action:
         toolbox.write_file(PurePosixPath("PROMPT"), prompt)