feat: add docstring prompt (#26)

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/bots/__init__.py

@@ -19,6 +19,19 @@
 
 
 def load_bot(config: BotConfig) -> Bot:
+    """Load and return a Bot instance using the provided configuration.
+
+    If a pythonpath is specified in the config and not already present in
+    sys.path, it is added. The function expects the config.factory in the
+    format 'module:symbol' or 'symbol'. If only 'symbol' is provided, the
+    current module is used.
+
+    Args:
+        config: BotConfig object containing bot configuration details.
+
+    Raises:
+        NotImplementedError: If the specified factory cannot be found.
+    """
     if config.pythonpath and config.pythonpath not in sys.path:
         sys.path.insert(0, config.pythonpath)
 
@@ -39,6 +52,14 @@ def load_bot(config: BotConfig) -> Bot:
 
 
 def openai_bot(**kwargs) -> Bot:
+    """Instantiate and return an OpenAIBot with provided keyword arguments.
+
+    This function imports the OpenAIBot class from the openai module and
+    returns an instance configured with the provided arguments.
+
+    Args:
+        **kwargs: Arbitrary keyword arguments used to configure the bot.
+    """
     from .openai import OpenAIBot
 
     return OpenAIBot(**kwargs)

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,30 +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(
-            loader=jinja2.FileSystemLoader(
-                [Config.folder_path() / "prompts", str(_prompt_root)]
-            ),
-            autoescape=False,
-            keep_trailing_newline=True,
-            auto_reload=False,
-        )
-        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"]
 
 
@@ -168,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,53 @@
+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:
+        """Parse arguments into a TemplatedPrompt
+
+        Args:
+            name: The name of the template.
+            *args: Additional arguments for context, expected in 'key=value'
+                format.
+        """
+        return cls(name, dict(e.split("=", 1) for e in args))
+
+
+class PromptRenderer:
+    """Renderer for prompt templates using Jinja2"""
+
+    def __init__(self, env: jinja2.Environment) -> None:
+        self._environment = env
+
+    @classmethod
+    def for_repo(cls, repo: git.Repo) -> Self:
+        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().splitlines(),
+        }
+        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

@@ -0,0 +1,37 @@
+{% 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
+which could be surprising, and make sure to mention those in the docstring.
+
+Docstrings should use the "Args" format for arguments. See the following
+examples:
+
+
+```python
+def write_file(path: Path, contents: str) -> None:
+    """Updates a file's contents
+
+    Args:
+        path: Path to the file to update.
+        contents: New file contents.
+    """
+  ...
+
+class Renderer:
+    """A smart renderer"""
+
+    ...
+```
+
+Additionally, the first paragraph of each docstring should fit in a single line
+and not include a period at the end. It should be a brief summary of the
+symbol's functionality.
+
+{% include "includes/file-list.jinja" %}

src/git_draft/prompts/add-test.jinja

@@ -1 +1,11 @@
-Add a test for {{ symbol }}.
+Add tests for {{ symbol }}. Follow existing conventions when implementing the
+tests. For example, if the surrounding code uses fixtures, do so as well.
+
+{% if src_path is defined %}The symbol to be tested is defined in {{ src_path }}.{% endif %}
+
+{% if test_path is defined %}The tests should be added to {{ test_path }}.{% endif %}
+
+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.
+
+{% 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