Extract InheritanceRef value object from InheritanceLoader.get_source

project_forge/rendering/environment.py

@@ -2,6 +2,7 @@
 
 import logging
 import re
+from dataclasses import dataclass
 from typing import Any, Callable, Optional
 
 from jinja2 import BaseLoader, Environment, TemplateNotFound, Undefined
@@ -11,6 +12,30 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class InheritanceRef:
+    """A parsed reference to one rung of a template's inheritance chain.
+
+    `InheritanceLoader` addresses a specific ancestor of a template with an
+    `"{index}/{name}"` string (e.g. `"2/a.txt"`); this is the value object for that string.
+    """
+
+    name: str
+    index: int = 0
+
+    @classmethod
+    def parse(cls, template: str) -> "InheritanceRef":
+        """Parse a loader template name, e.g. `"2/a.txt"` or the un-prefixed `"a.txt"`."""
+        bits = template.split("/", maxsplit=1)
+        if len(bits) == 2 and bits[0].isdigit():
+            return cls(name=bits[1], index=int(bits[0]))
+        return cls(name=template)
+
+    def format(self) -> str:
+        """Render this reference back into the loader's `"{index}/{name}"` string form."""
+        return f"{self.index}/{self.name}"
+
+
 class SuperUndefined(Undefined):
     """Let calls to super() work ok."""
 
@@ -35,33 +60,26 @@ def __init__(self, inheritance_map: InheritanceMap):
 
     def get_source(self, environment: Environment, template: str) -> tuple[str, str | None, Callable[[], bool] | None]:
         """Load the template."""
-        # Parse the name of the template
-        bits = template.split("/", maxsplit=1)
-        index = 0
-        if len(bits) == 2 and bits[0].isdigit():
-            index = int(bits[0])
-            template_name = bits[1]
-        else:
-            template_name = template
+        ref = InheritanceRef.parse(template)
 
         # Get template inheritance
-        inheritance = self.templates.inheritance(template_name)
+        inheritance = self.templates.inheritance(ref.name)
         inheritance_len = len(inheritance)
 
         if not inheritance:
-            raise TemplateNotFound(template_name)
+            raise TemplateNotFound(ref.name)
 
         # Load the template from the index
-        if index >= inheritance_len:
+        if ref.index >= inheritance_len:
             raise TemplateNotFound(template)  # Maybe this wasn't one of our customized extended paths
 
-        template_file = inheritance[index]
+        template_file = inheritance[ref.index]
 
         if not template_file.is_renderable:
             raise TemplateNotFound(template)
 
         path = template_file.path
-        logger.debug(f"Loading template {template_name} from: {path}")
+        logger.debug(f"Loading template {ref.name} from: {path}")
         source = path.read_text()
 
         # look for an `extends` tag
@@ -71,12 +89,13 @@ def get_source(self, environment: Environment, template: str) -> tuple[str, str
             self.extends_re.format(block_start_string=block_start_string, block_end_string=block_end_string)
         )
         if match := regex.search(source):
-            if index == len(inheritance) - 1:
+            if ref.index == inheritance_len - 1:
                 # we've reached our last template, so we must remove the `extends` tag completely
                 source = source.replace(match[0], "")
             else:
                 # rewrite the `extends` tag to reference the next item in the inheritance
-                source = source.replace(match[1], f"{index + 1}/{match[1]}")
+                extended_ref = InheritanceRef(name=match[1], index=ref.index + 1)
+                source = source.replace(match[1], extended_ref.format())
 
         return source, None, lambda: True
 

tests/test_rendering/test_environment.py

@@ -6,12 +6,28 @@
 from jinja2 import Environment, TemplateNotFound
 from jinja2.exceptions import UndefinedError
 
-from project_forge.rendering.environment import InheritanceLoader, SuperUndefined
+from project_forge.rendering.environment import InheritanceLoader, InheritanceRef, SuperUndefined
 from project_forge.rendering.templates import InheritanceMap, ProcessMode, TemplateFile
 
 RW_MODE = ProcessMode.render | ProcessMode.write
 
 
+class TestInheritanceRef:
+    """Tests for InheritanceRef."""
+
+    def test_parse_unprefixed_name_defaults_to_index_zero(self):
+        """Parsing a plain template name yields index 0."""
+        assert InheritanceRef.parse("a.txt") == InheritanceRef(name="a.txt", index=0)
+
+    def test_parse_prefixed_name_extracts_index(self):
+        """Parsing a `"N/name"` string extracts the index and name."""
+        assert InheritanceRef.parse("2/a.txt") == InheritanceRef(name="a.txt", index=2)
+
+    def test_format_renders_prefixed_string(self):
+        """`format()` renders the `"N/name"` string form."""
+        assert InheritanceRef(name="a.txt", index=2).format() == "2/a.txt"
+
+
 @pytest.fixture
 def init_map(tmp_path: Path) -> InheritanceMap:
     """Create a simple template inheritance map."""
@@ -124,3 +140,25 @@ def test_get_source_for_last_template_removes_extends_tag(self, init_map: Inheri
 
         source = env.loader.get_source(env, f"2/{template_name}")
         assert source[0] == ""
+
+    def test_get_source_rewrites_extends_using_the_captured_tag_text(self, tmp_path: Path):
+        """The rewritten `extends` reference must reuse the text actually written in the tag.
+
+        A template's `{% extends %}` string need not match the key it was cataloged under
+        (e.g. it may reference a path prefix); the rewrite must not substitute the map key instead.
+        """
+        p1 = tmp_path / "dir1" / "a.txt"
+        p1.parent.mkdir(parents=True, exist_ok=True)
+        p1.write_text("{% extends 'templates/a.txt' %}")
+        p2 = tmp_path / "dir2" / "a.txt"
+        p2.parent.mkdir(parents=True, exist_ok=True)
+        p2.write_text("done")
+        inheritance_map = InheritanceMap(
+            {"a.txt": TemplateFile(p2, RW_MODE)},
+            {"a.txt": TemplateFile(p1, RW_MODE)},
+        )
+        loader = InheritanceLoader(inheritance_map)
+        env = Environment(loader=loader)
+
+        source, _, _ = env.loader.get_source(env, "a.txt")
+        assert "1/templates/a.txt" in source