feat: add extra and template_name fields to generator configs

Author: esynr3z

schemas/corsair-build-schema.json

@@ -15,6 +15,12 @@
       "additionalProperties": false,
       "description": "Configuration for the Markdown generator.",
       "properties": {
+        "extra": {
+          "additionalProperties": true,
+          "description": "Extra configuration parameters for the generator.",
+          "title": "Extra",
+          "type": "object"
+        },
         "kind": {
           "const": "markdown",
           "default": "markdown",
@@ -34,6 +40,12 @@
           "title": "Title",
           "type": "string"
         },
+        "template_name": {
+          "default": "regmap.md.j2",
+          "description": "Name of the Jinja2 template to use.",
+          "title": "Template Name",
+          "type": "string"
+        },
         "print_images": {
           "default": false,
           "description": "Enable generating images for bit fields of a register.",
@@ -56,6 +68,7 @@
         "wavedrom": {
           "$ref": "#/$defs/corsair___generators__wavedrom__WaveDromGenerator__Config",
           "default": {
+            "extra": {},
             "kind": "wavedrom",
             "vspace": 80,
             "hspace": 800,
@@ -79,6 +92,12 @@
       "additionalProperties": false,
       "description": "Configuration for the Verilog generator.",
       "properties": {
+        "extra": {
+          "additionalProperties": true,
+          "description": "Extra configuration parameters for the generator.",
+          "title": "Extra",
+          "type": "object"
+        },
         "kind": {
           "const": "verilog",
           "default": "verilog",
@@ -98,6 +117,12 @@
       "additionalProperties": false,
       "description": "Configuration for the VHDL generator.",
       "properties": {
+        "extra": {
+          "additionalProperties": true,
+          "description": "Extra configuration parameters for the generator.",
+          "title": "Extra",
+          "type": "object"
+        },
         "kind": {
           "const": "vhdl",
           "default": "vhdl",
@@ -117,6 +142,12 @@
       "additionalProperties": false,
       "description": "Configuration for the WaveDrom generator.",
       "properties": {
+        "extra": {
+          "additionalProperties": true,
+          "description": "Extra configuration parameters for the generator.",
+          "title": "Extra",
+          "type": "object"
+        },
         "kind": {
           "const": "wavedrom",
           "default": "wavedrom",

src/corsair/_generators/base.py

@@ -15,7 +15,7 @@
     from corsair._model import Map
 
 import jinja2
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 
 from corsair._templates import TemplateEnvironment
 
@@ -81,6 +81,9 @@ def __str__(self) -> str:
 class GeneratorConfig(BaseModel, ABC):
     """Base configuration for a generator."""
 
+    extra: dict[str, Any] = Field(default_factory=dict)
+    """Extra configuration parameters for the generator."""
+
     model_config = ConfigDict(
         extra="forbid",
         use_attribute_docstrings=True,
@@ -136,7 +139,7 @@ def __call__(self) -> TypeGenerator[Path, None, None]:
 
     def _render_to_text(self, template_name: str, context: dict[str, Any]) -> str:
         """Render text with Jinja2."""
-        env = TemplateEnvironment(searchpaths=self.template_searchpaths)
+        env = TemplateEnvironment(searchpath=self.template_searchpaths)
         template = env.get_template(template_name)
         return template.render(context)
 

src/corsair/_generators/markdown.py

@@ -15,8 +15,6 @@
 class MarkdownGenerator(Generator):
     """Markdown file generator for a register map."""
 
-    template_name: str = "regmap.md.j2"
-
     class Config(GeneratorConfig):
         """Configuration for the Markdown generator."""
 
@@ -29,6 +27,9 @@ class Config(GeneratorConfig):
         title: str = "Register Map"
         """Document title."""
 
+        template_name: str = "regmap.md.j2"
+        """Name of the Jinja2 template to use."""
+
         print_images: bool = False
         """Enable generating images for bit fields of a register."""
 
@@ -65,7 +66,7 @@ def _generate(self) -> TypeGenerator[Path, None, None]:
         }
 
         yield self._render_to_file(
-            template_name=self.template_name,
+            template_name=self.config.template_name,
             context=context,
             file_name=self.config.file_name,
         )

src/corsair/_templates/env.py

@@ -5,14 +5,47 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-from jinja2 import Environment, FileSystemLoader, StrictUndefined
+from jinja2 import Environment, FileSystemLoader, StrictUndefined, TemplateNotFound
 
 from corsair._version import VERSION
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from typing_extensions import Self
 
 
+class _EnhancedFileSystemLoader(FileSystemLoader):
+    """Custom Jinja2 FileSystemLoader that supports both relative and absolute paths."""
+
+    def __init__(self, searchpath: list[Path], **kwargs: Any) -> None:
+        super().__init__(searchpath=searchpath, **kwargs)
+
+    def get_source(
+        self,
+        environment: Environment,
+        template: str,
+    ) -> tuple[str, str, Callable[[], bool]]:
+        """Get the template source, filename and reload helper for a template.
+
+        Checks for absolute paths first, otherwise delegates to the parent FileSystemLoader.
+        """
+        template_path = Path(template)
+        if template_path.is_absolute():
+            if not template_path.is_file():
+                raise TemplateNotFound(template)
+
+            mtime = template_path.stat().st_mtime
+            with template_path.open("r", encoding="utf-8") as f:
+                source = f.read()
+
+            # Return the source, filename (absolute path), and a mtime check function
+            return source, str(template_path), lambda: template_path.stat().st_mtime == mtime
+
+        # If it's not an absolute path, delegate to the parent class method
+        return super().get_source(environment, template)
+
+
 class TemplateEnvironment(Environment):
     """Singleton environment for managing Jinja2 templates."""
 
@@ -24,14 +57,16 @@ def __new__(cls, *args: Any, **kwargs: Any) -> Self:  # noqa: ARG003
             cls._instance = super().__new__(cls)
         return cls._instance
 
-    def __init__(self, searchpaths: list[Path] | None = None) -> None:
+    def __init__(self, searchpath: list[Path] | None = None) -> None:
         """Initialize the template environment."""
-        if not searchpaths:
-            searchpaths = []
-        searchpaths.append(Path(__file__).parent)
+        if not searchpath:
+            searchpath = []
+        # Always include the directory containing the built-in templates
+        if Path(__file__).parent not in searchpath:
+            searchpath.append(Path(__file__).parent)
 
         super().__init__(
-            loader=FileSystemLoader(searchpath=searchpaths),
+            loader=_EnhancedFileSystemLoader(searchpath=searchpath),
             trim_blocks=True,
             lstrip_blocks=True,
             undefined=StrictUndefined,  # to throw exception on any undefined variable within template

tests/generators/test_markdown.py

@@ -213,3 +213,45 @@ def test_wavedrom_dump_json(tmp_path: Path, simple_regmap: csr.Map) -> None:
     assert expected_data_dir.is_dir()
     assert {f for f in expected_data_dir.iterdir() if f.suffix == ".svg"} == expected_svg_files
     assert {f for f in expected_data_dir.iterdir() if f.suffix == ".json"} == expected_json_files
+
+
+def test_custom_template_and_extra_param(tmp_path: Path, simple_regmap: csr.Map) -> None:
+    """Test generation with a custom template and extra config parameters."""
+    custom_template_name = "custom_template.md.j2"
+    custom_template_content = """
+# Custom Template Test
+
+Registers:
+{% for item in regmap.items %}
+- {{ item.name }}
+{% endfor %}
+
+Custom Param: {{ cfg.extra['my_param'] }}
+"""
+    custom_template_file = tmp_path / custom_template_name
+    custom_template_file.write_text(custom_template_content)
+
+    custom_key = "my_param"
+    custom_value = "hello_world"
+    config = csr.MarkdownGenerator.Config(template_name=str(custom_template_file), extra={custom_key: custom_value})
+
+    gen = csr.MarkdownGenerator(
+        label="test_md_gen_custom_tmpl",
+        register_map=simple_regmap,
+        config=config,
+        output_dir=tmp_path,
+    )
+    generated_files = list(gen())
+
+    # Check that the default file name is used if not specified
+    expected_file = tmp_path / "regmap.md"
+    assert generated_files == [expected_file]
+    assert expected_file.is_file()
+
+    # Check the content of the generated file
+    content = expected_file.read_text()
+    assert "# Custom Template Test" in content
+    assert "Registers:" in content
+    assert "- reg1" in content
+    assert "- reg2" in content
+    assert f"Custom Param: {custom_value}" in content