Repository: pygfx/pygfx
Pull request #1083 (2025-05-06 09:19:51)

Properly implement `register_wgsl_loader`
Ref #1077

We implemented `register_wgsl_loader()` with the idea to let downstream code register their own loader. It seems like the implementation was a stub that was never properly finished: there is no test, we don't use it ourselves, and it's simply broken.

This PR fixes that, allowing downstream code to provide either a `jinja2.BaseLoader`, a function, or a dict.



maharshi-gor: Thank you @almarklein for the work this will be really helpful.

Patch:
diff --git a/pygfx/renderers/wgpu/shader/templating.py b/pygfx/renderers/wgpu/shader/templating.py
index dbfe7d0..e60458a 100644
--- a/pygfx/renderers/wgpu/shader/templating.py
+++ b/pygfx/renderers/wgpu/shader/templating.py
@@ -1,7 +1,6 @@
 import jinja2
 
-loader = jinja2.PrefixLoader({}, delimiter=".")
-loader.mapping["pygfx"] = jinja2.PackageLoader("pygfx.renderers.wgpu.wgsl", ".")
+root_loader = jinja2.PrefixLoader({}, delimiter=".")
 
 jinja_env = jinja2.Environment(
     block_start_string="{$",
@@ -10,11 +9,11 @@ jinja_env = jinja2.Environment(
     variable_end_string="}}",
     line_statement_prefix="$$",
     undefined=jinja2.StrictUndefined,
-    loader=loader,
+    loader=root_loader,
 )
 
 
-def register_wgsl_loader(context, func):
+def register_wgsl_loader(context, loader):
     """Register a source for shader snippets.
 
     When code is encountered that looks like::
@@ -22,23 +21,33 @@ def register_wgsl_loader(context, func):
        {$ include 'some_context.name.wgsl' $}
 
     The loader for "some_context" is looked up and used to load the wgsl to include.
-    This function allows registering a loader for your downstream package.
+    This function allows registering a loader for your downstream package or application.
 
     Parameters
     ----------
     context : str
         The context of the loader.
-    func: callable
-        The function that will be called when a shader is loaded for the given context.
-        The function must accept one positional argument (the name to include).
+    loader: jinja2.BaseLoader | callable | dict
+        The loader to use for this context. If a function is given, it must accept one
+        positional argument (the name to include).
     """
     if not (isinstance(context, str) and "." not in context):
         raise TypeError("Wgsl load context must be a string witout dots.")
-    if not callable(func):
-        raise TypeError("The given wgsl load func must be callable.")
-    if context in loader.mapping:
+    if context in root_loader.mapping:
         raise RuntimeError(f"A loader is already registered for '{context}'.")
-    loader.mapping[context] = func
+    if isinstance(loader, jinja2.BaseLoader):
+        root_loader.mapping[context] = loader
+    elif isinstance(loader, dict):
+        root_loader.mapping[context] = jinja2.DictLoader(loader)
+    elif callable(loader):
+        root_loader.mapping[context] = jinja2.FunctionLoader(loader)
+    else:
+        raise TypeError(
+            f"The given wgsl loader must be a jinja2.BaseLoader, function, or dict. Not {loader!r}"
+        )
+
+
+register_wgsl_loader("pygfx", jinja2.PackageLoader("pygfx.renderers.wgpu.wgsl", "."))
 
 
 def apply_templating(code, **kwargs):


Test patch:
diff --git a/tests/renderers/test_shader_composer.py b/tests/renderers/test_shader_composer.py
index 7ddfb34..4ab2022 100644
--- a/tests/renderers/test_shader_composer.py
+++ b/tests/renderers/test_shader_composer.py
@@ -1,8 +1,10 @@
+import pygfx
 from pygfx.renderers.wgpu.shader import BaseShader
-from pygfx.renderers.wgpu import Binding
+from pygfx.renderers.wgpu import Binding, register_wgsl_loader
 from pygfx.utils import array_from_shadertype
 from pygfx import Buffer
 from pytest import raises
+import jinja2
 import numpy as np
 
 
@@ -187,7 +189,43 @@ def test_uniform_definitions():
         shader.define_binding(0, 0, Binding("zz", "buffer/uniform", struct))
 
 
+def test_custom_wgsl_loaders():
+    class MyLoader(jinja2.BaseLoader):
+        def get_source(self, environment, x):
+            return f"// loader: {x}", None, None
+
+    my_loader = MyLoader()
+
+    def my_func(x):
+        return f"// func: {x}"
+
+    my_dict = {"foo.wgsl": "// dict: foo.wgsl"}
+
+    register_wgsl_loader("testloader1", my_loader)
+    register_wgsl_loader("testloader2", my_func)
+    register_wgsl_loader("testloader3", my_dict)
+
+    with raises(RuntimeError):
+        register_wgsl_loader("testloader3", {})
+
+    code1 = """
+    {$ include 'testloader1.foo.wgsl' $}
+    {$ include 'testloader2.foo.wgsl' $}
+    {$ include 'testloader3.foo.wgsl' $}
+    """
+
+    code2 = """
+    // loader: foo.wgsl
+    // func: foo.wgsl
+    // dict: foo.wgsl
+    """
+
+    code3 = pygfx.renderers.wgpu.shader.templating.apply_templating(code1)
+    assert code3 == code2
+
+
 if __name__ == "__main__":
     test_templating()
     test_logic_beyond_templating()
     test_uniform_definitions()
+    test_custom_wgsl_loaders()