fix(loader): Add loaded module to sys.modules to help libraries using data from annotation scopes

PR-12: #12

Author: noirbizarre

src/copier_templates_extensions/extensions/loader.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import sys
 from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -80,9 +81,10 @@ def _import_template_module(self, relative_path: str | Path) -> ModuleType:
                 "Please report this issue to the template maintainers.",
             )
         spec = spec_from_file_location(
-            f"copier_templates_extensions.{module_name}",
+            module_full_name := f"copier_templates_extensions.{module_name}",
             template_relative_path,
         )
         module = module_from_spec(spec)  # type: ignore[arg-type]
+        sys.modules[module_full_name] = module
         spec.loader.exec_module(module)  # type: ignore[union-attr]
         return module

tests/fixtures/dataclass/copier.yml

@@ -0,0 +1,4 @@
+_subdirectory: template
+_jinja_extensions:
+  - copier_templates_extensions.TemplateExtensionLoader
+  - extensions/ctx.py:Success

tests/fixtures/dataclass/extensions/ctx.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from jinja2.ext import Extension
+
+@dataclass
+class Test:
+    key: str
+    value: str
+
+class Success(Extension):
+    def __init__(self, environment):
+        super().__init__(environment)
+        environment.globals.update(success=True)

tests/fixtures/dataclass/template/result.txt.jinja

@@ -0,0 +1 @@
+Success variable: {{ success|default("not set") }}

tests/test_extensions.py

@@ -18,6 +18,7 @@
         ("modifying_context", "True"),
         ("not_updating_context_for_prompts", "True"),
         ("loading_normal_extension", "not set"),
+        ("dataclass", "True"),
     ],
 )
 def test_extensions(tmp_path: Path, template_name: str, expected: str) -> None: