Let custom generators declare their own source includes

README.md

@@ -274,8 +274,33 @@ r = Rectangle(4, 5)
   and a type name that is defined in more than one header is treated as ambiguous
   and left for a manual `source_includes`. It is a no-op with
   `common_include_file`, and does not cover types that appear only in
-  hand-written `custom_generator` code (cppwg cannot see those). See the
-  `MeshFactory` class in `examples/cells/dynamic/config.yaml`.
+  hand-written `custom_generator` code (cppwg cannot see those) — a generator
+  that names such types can add their headers itself by overriding
+  `get_source_includes()` (see below). See the `MeshFactory` class in
+  `examples/cells/dynamic/config.yaml`.
+- A `custom_generator` can emit code that references types cppwg never sees in
+  the parsed signatures (e.g. a template-method instantiation like
+  `AddCellWriter<CellAgesWriter>` built from a hard-coded list). Those headers
+  cannot be auto-included and would otherwise have to be repeated under
+  `source_includes`. Instead, override `get_source_includes()` on the generator
+  (a subclass of `cppwg.templates.custom.Custom`) to return the header names —
+  cppwg adds them to the wrapper's `#include` block (deduplicated, quoted or
+  `<...>` form), so the generator and the includes its code needs live in one
+  place:
+
+  ```python
+  class PopulationWriterCustomTemplate(cppwg.templates.custom.Custom):
+      WRITERS = ["CellAgesWriter", "CellIdWriter", ...]
+
+      def get_class_cpp_def_code(self, class_name):
+          return "".join(
+              f'.def("AddCellWriter{w}", &{class_name}::AddCellWriter<{w}>)\n'
+              for w in self.WRITERS
+          )
+
+      def get_source_includes(self, *args, **kwargs):
+          return [f"{w}.hpp" for w in self.WRITERS]
+  ```
 - A wrapped class can inherit from a base class wrapped in a **different
   module**, as long as that base is registered somewhere that gets imported
   first. Opt in per module with `imports`, which lists the Python modules to

cppwg/templates/custom.py

@@ -25,6 +25,19 @@ def get_class_cpp_def_code(self, *args, **kwargs) -> str:
 
         return ""
 
+    def get_source_includes(self, *args, **kwargs) -> list:
+        """
+        Return headers to add to the #include block of the class wrapper.
+
+        Use this for headers the generated code needs but that cppwg cannot
+        infer from the parsed C++ signatures - e.g. types named only inside the
+        get_class_cpp_def_code() output, which auto-include detection never sees.
+        Each entry is a header spelled as under source_includes: quoted
+        ("Foo.hpp") or angle-bracket (<foo>). Returns an empty list by default.
+        """
+
+        return []
+
     def get_module_pre_code(self) -> str:
         """
         Return a string of C++ code to be inserted before the module

cppwg/utils/utils.py

@@ -40,6 +40,37 @@ def ensure_trailing_newline(code: str) -> str:
     return code
 
 
+def call_generator_hook(
+    generator: Any, method_name: str, default: Any, *args: Any
+) -> Any:
+    """
+    Call an optional custom-generator hook, returning ``default`` if absent.
+
+    A custom generator need not subclass ``cppwg.templates.custom.Custom`` or
+    implement every hook. A missing or non-callable hook - including the case of
+    no generator at all (``generator`` is None) - yields ``default`` rather than
+    raising, so partial or legacy generators keep working.
+
+    Parameters
+    ----------
+    generator : Any
+        The custom generator instance, or None.
+    method_name : str
+        The name of the hook to call, e.g. "get_class_cpp_def_code".
+    default : Any
+        The value to return when the hook is absent.
+    *args : Any
+        Positional arguments passed to the hook.
+
+    Returns
+    -------
+    Any
+        The hook's return value, or ``default`` if the hook is absent.
+    """
+    hook = getattr(generator, method_name, None)
+    return hook(*args) if callable(hook) else default
+
+
 def write_file_if_changed(filepath: str, content: str, overwrite: bool = False) -> bool:
     """
     Write content to filepath unless an identical file already exists.

cppwg/writers/class_writer.py

@@ -13,6 +13,7 @@
     CPPWG_HEADER_COLLECTION_FILENAME,
 )
 from cppwg.utils.utils import (
+    call_generator_hook,
     ensure_trailing_newline,
     type_string_matches,
     write_file_if_changed,
@@ -177,6 +178,16 @@ def add(line: str) -> None:
         ):
             add(self._format_include(source_include))
 
+        # Headers a custom generator's emitted code needs. A generator can name
+        # types in its get_class_cpp_def_code() output that never appear in the
+        # parsed signatures (e.g. AddCellWriter<CellAgesWriter>), so cppwg cannot
+        # auto-detect them; it declares them via the optional get_source_includes()
+        # hook (see call_generator_hook).
+        for header in call_generator_hook(
+            self.class_info.custom_generator_instance, "get_source_includes", []
+        ):
+            add(self._format_include(header))
+
         source_file = self.class_info.source_file
         if not source_file:
             source_file = os.path.basename(self.class_info.decls[0].location.file_name)
@@ -487,8 +498,8 @@ def build_class_register(self, template_idx: int) -> tuple[str, str]:
         )
 
         block = self.wrapper_templates["class_cpp_register"].substitute(
-            generator_pre_code=(
-                generator.get_class_cpp_pre_code(class_py_name) if generator else ""
+            generator_pre_code=call_generator_hook(
+                generator, "get_class_cpp_pre_code", "", class_py_name
             ),
             class_py_name=class_py_name,
             class_cpp_name=class_cpp_name,
@@ -503,7 +514,9 @@ def build_class_register(self, template_idx: int) -> tuple[str, str]:
             # missing newline (or a trailing // comment) could swallow the
             # statement terminator. See ensure_trailing_newline.
             generator_def_code=ensure_trailing_newline(
-                generator.get_class_cpp_def_code(class_py_name) if generator else ""
+                call_generator_hook(
+                    generator, "get_class_cpp_def_code", "", class_py_name
+                )
             ),
             suffix_code=self.suffix_code(),
         )
@@ -549,8 +562,8 @@ def build_struct_enum_register(self, template_idx: int) -> str:
         )
 
         return self.wrapper_templates["struct_enum_register"].substitute(
-            generator_pre_code=(
-                generator.get_class_cpp_pre_code(class_py_name) if generator else ""
+            generator_pre_code=call_generator_hook(
+                generator, "get_class_cpp_pre_code", "", class_py_name
             ),
             class_py_name=class_py_name,
             class_cpp_name=class_cpp_name,

cppwg/writers/module_writer.py

@@ -5,7 +5,11 @@
 from typing import TYPE_CHECKING
 
 from cppwg.utils.constants import CPPWG_EXT, CPPWG_HEADER_COLLECTION_FILENAME
-from cppwg.utils.utils import ensure_trailing_newline, write_file_if_changed
+from cppwg.utils.utils import (
+    call_generator_hook,
+    ensure_trailing_newline,
+    write_file_if_changed,
+)
 from cppwg.writers.class_writer import CppClassWrapperWriter
 from cppwg.writers.free_function_writer import CppFreeFunctionWrapperWriter
 
@@ -217,7 +221,7 @@ def build_module_context(self) -> dict[str, str]:
             # closing brace, so normalise both to end with a newline (see
             # ensure_trailing_newline) to avoid producing invalid C++.
             "module_pre_code": ensure_trailing_newline(
-                generator.get_module_pre_code() if generator else ""
+                call_generator_hook(generator, "get_module_pre_code", "")
             ),
             "class_includes": class_includes,
             "full_module_name": self.full_module_name,
@@ -228,7 +232,7 @@ def build_module_context(self) -> dict[str, str]:
             "free_functions": free_functions,
             "register_calls": register_calls,
             "module_code": ensure_trailing_newline(
-                generator.get_module_code() if generator else ""
+                call_generator_hook(generator, "get_module_code", "")
             ),
         }
 

tests/test_class_writer.py

@@ -282,6 +282,9 @@ def get_class_cpp_pre_code(self, class_py_name):
     def get_class_cpp_def_code(self, class_py_name):
         return ""
 
+    def get_source_includes(self, *args, **kwargs):
+        return []
+
 
 def test_generator_pre_code_follows_alias_typedef():
     """A custom generator's pre-code is emitted after the alias typedef.
@@ -553,6 +556,100 @@ def test_includes_block_dedups_auto_include_with_source_include():
     )
 
 
+class _SourceIncludeGen:
+    """Custom generator that declares headers its generated code needs."""
+
+    def __init__(self, headers):
+        self._headers = headers
+
+    def get_source_includes(self, *args, **kwargs):
+        return self._headers
+
+
+def test_includes_block_emits_generator_source_includes():
+    """A custom generator's get_source_includes() headers are added to the block.
+
+    Covers the category the auto-include detection cannot see: types named only
+    in the generator's emitted code (e.g. AddCellWriter<CellAgesWriter>), whose
+    headers the generator supplies itself. Angle-bracket and quoted forms both
+    work.
+    """
+    class_info = _FakeClassInfo(
+        "Population",
+        object(),
+        {"common_include_file": False},
+        "Population.hpp",
+        generator=_SourceIncludeGen(["CellAgesWriter.hpp", "<memory>"]),
+    )
+    writer = _make_writer(class_info)
+
+    assert writer.includes_block() == (
+        '#include "CellAgesWriter.hpp"\n'
+        "#include <memory>\n"
+        '#include "Population.hpp"\n'
+    )
+
+
+def test_includes_block_dedups_generator_and_source_includes():
+    """A header from both source_includes and the generator emits once."""
+    class_info = _FakeClassInfo(
+        "Foo",
+        object(),
+        {"common_include_file": False, "source_includes": ["Shared.hpp"]},
+        "Foo.hpp",
+        generator=_SourceIncludeGen(["Shared.hpp"]),
+    )
+    writer = _make_writer(class_info)
+
+    assert writer.includes_block() == (
+        '#include "Shared.hpp"\n'  # once (source_includes, then generator deduped)
+        '#include "Foo.hpp"\n'
+    )
+
+
+def test_includes_block_generator_without_source_includes_hook():
+    """A legacy generator lacking get_source_includes() does not break generation.
+
+    Generators that only implement the pre/def-code methods (and do not subclass
+    Custom) must keep working - the optional hook is skipped, not required.
+    """
+
+    class LegacyGen:
+        def get_class_cpp_pre_code(self, *args):
+            return ""
+
+        def get_class_cpp_def_code(self, *args):
+            return ""
+
+    class_info = _FakeClassInfo(
+        "Foo",
+        object(),
+        {"common_include_file": False},
+        "Foo.hpp",
+        generator=LegacyGen(),
+    )
+    writer = _make_writer(class_info)
+
+    # No AttributeError; just the class's own header.
+    assert writer.includes_block() == '#include "Foo.hpp"\n'
+
+
+def test_includes_block_generator_source_includes_skipped_in_common():
+    """In common-include-file mode, generator includes are moot (collection has all)."""
+    class_info = _FakeClassInfo(
+        "Foo",
+        object(),
+        {"common_include_file": True},
+        "Foo.hpp",
+        generator=_SourceIncludeGen(["Writer.hpp"]),
+    )
+    writer = _make_writer(class_info)
+
+    assert writer.includes_block() == (
+        '#include "wrapper_header_collection.cppwg.hpp"\n'
+    )
+
+
 def test_includes_block_emits_typecasters_common():
     """Detected caster headers follow the header collection in common mode."""
     class_info = _FakeClassInfo(

tests/test_utils.py

@@ -5,6 +5,7 @@
 import pytest
 
 from cppwg.utils.utils import (
+    call_generator_hook,
     ensure_trailing_newline,
     find_template_instantiations_in_source,
     find_template_instantiations_in_source_file,
@@ -350,3 +351,33 @@ def test_overwrite_forces_rewrite_when_unchanged(tmp_path):
 
     assert wrote is True
     assert os.stat(filepath).st_mtime_ns != old_time_ns
+
+
+def test_call_generator_hook():
+    """An optional generator hook is called if present, else the default is used."""
+
+    class PartialGenerator:
+        def get_module_code(self):
+            return "// code"
+
+        def get_class_cpp_pre_code(self, class_py_name):
+            return f"// {class_py_name}"
+
+    gen = PartialGenerator()
+
+    # Present hooks are called (with args forwarded).
+    assert call_generator_hook(gen, "get_module_code", "") == "// code"
+    assert call_generator_hook(gen, "get_class_cpp_pre_code", "", "Foo") == "// Foo"
+
+    # A missing hook (generator does not subclass Custom / omits it) -> default.
+    assert call_generator_hook(gen, "get_source_includes", []) == []
+    assert call_generator_hook(gen, "get_class_cpp_def_code", "", "Foo") == ""
+
+    # No generator at all -> default.
+    assert call_generator_hook(None, "get_module_code", "DEFAULT") == "DEFAULT"
+
+    # A non-callable attribute of that name -> default (treated as absent).
+    class Weird:
+        get_module_code = "not callable"
+
+    assert call_generator_hook(Weird(), "get_module_code", "DEFAULT") == "DEFAULT"