feat: custom extension support simply pass in code

jinja2_test.go

@@ -125,6 +125,21 @@ def multiply(a, b):
 `))
 	assert.NoError(t, err)
 	assert.Equal(t, "test - 6", s)
+
+	s, err = j2.RenderString("test - {{ test_var1 | minus(-2) }}", WithExtension(`
+from jinja2.ext import Extension
+
+class DemoExtension(Extension):
+    def __init__(self, environment):
+        super().__init__(environment)
+        environment.filters['minus'] = self.minus
+
+    @staticmethod
+    def minus(a, b):
+        return int(a) - int(b)
+`))
+	assert.NoError(t, err)
+	assert.Equal(t, "test - 3", s)
 }
 
 type testStruct struct {

opts.go

@@ -108,6 +108,10 @@ func WithFilter(name string, code string) Jinja2Opt {
 	}
 }
 
+// WithExtension adds a custom extension to the engine
+//
+// You can pass in an import path to an extension (for using some built-in extension, e.g. `jinja2.ext.debug`)
+// Or you can pass in the code of an extension (there must be exactly one class inherited from `jinja2.ext.Extension`)
 func WithExtension(e string) Jinja2Opt {
 	return func(o *jinja2Options) {
 		o.Extensions = append(o.Extensions, e)

python_src/go_jinja2/jinja2_renderer.py

@@ -1,4 +1,5 @@
 from jinja2 import StrictUndefined, ChainableUndefined, ChoiceLoader
+from jinja2.ext import Extension
 
 from .jinja2_utils import MyEnvironment, extract_template_error, RootTemplateLoader, SearchPathAbsLoader, \
     MyFileSystemLoader
@@ -42,7 +43,20 @@ def build_env(self):
         environment.globals.update(self.opts.get("globals", {}))
 
         for e in self.opts.get("extensions", []):
-            environment.add_extension(e)
+            e = str(e)
+            if e.count("\n") == 0:  # single line, assume it's a module name
+                environment.add_extension(e)
+            else:  # multi line, assume it's python code
+                track = {}
+                exec(e, track)
+                k = None
+                for v in track.values():
+                    if isinstance(v, type) and v != Extension and issubclass(v, Extension):
+                        k = v
+                        break
+                if k is None:
+                    raise AttributeError("No Extension subclass found in code")
+                environment.add_extension(k)
 
         for name, code in self.opts.get("filters", {}).items():
             track = {}