From 2c20f8c93c425c627d7267fd7a775f9ba603083d Mon Sep 17 00:00:00 2001
From: Pavel Tiunov <pavel.tiunov@gmail.com>
Date: Mon, 2 Oct 2023 11:51:42 -0700
Subject: [PATCH] feat(python): Support `@template.function("name")`
 definitions

---
 .../python/cube/src/__init__.py               | 56 +++++++++++++++++--
 .../cubejs-backend-native/src/python/entry.rs | 16 +++++-
 .../test/templates/jinja-instance.py          |  4 +-
 3 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/packages/cubejs-backend-native/python/cube/src/__init__.py b/packages/cubejs-backend-native/python/cube/src/__init__.py
index 5acc1c62517fe..2533a2d8dca6f 100644
--- a/packages/cubejs-backend-native/python/cube/src/__init__.py
+++ b/packages/cubejs-backend-native/python/cube/src/__init__.py
@@ -140,9 +140,9 @@ class AttrRef:
     config: Configuration
     attribute: str
 
-    def __init__(self, config: Configuration, attribue: str):
+    def __init__(self, config: Configuration, attribute: str):
         self.config = config
-        self.attribute = attribue
+        self.attribute = attribute
 
     def __call__(self, func):
         if not callable(func):
@@ -163,21 +163,65 @@ class TemplateException(Exception):
     pass
 
 class TemplateContext:
-    def function(self, func):
+    functions: dict[str, Callable]
+
+    def __init__(self):
+        self.functions = {}
+
+    def add_function(self, name, func):
         if not callable(func):
             raise TemplateException("function registration must be used with functions, actual: '%s'" % type(func).__name__)
-    
-        return context_func(func)
 
-    def filter(self, func):
+        self.functions[name] = func
+
+    def add_filter(self, name, func):
         if not callable(func):
             raise TemplateException("function registration must be used with functions, actual: '%s'" % type(func).__name__)
 
         raise TemplateException("filter registration is not supported")
 
+    def function(self, func):
+        if isinstance(func, str):
+            return TemplateFunctionRef(self, func)
+
+        self.add_function(func.__name__, func)
+        return func
+
+    def filter(self, func):
+        if isinstance(func, str):
+            return TemplateFilterRef(self, func)
+
+        self.add_filter(func.__name__, func)
+        return func
+
     def variable(self, func):
         raise TemplateException("variable registration is not supported")
 
+class TemplateFunctionRef:
+    context: TemplateContext
+    attribute: str
+
+    def __init__(self, context: TemplateContext, attribute: str):
+        self.context = context
+        self.attribute = attribute
+
+    def __call__(self, func):
+        self.context.add_function(self.attribute, func)
+        return func
+
+
+class TemplateFilterRef:
+    context: TemplateContext
+    attribute: str
+
+    def __init__(self, context: TemplateContext, attribute: str):
+        self.context = context
+        self.attribute = attribute
+
+    def __call__(self, func):
+        self.context.add_filter(self.attribute, func)
+        return func
+
 def context_func(func):
     func.cube_context_func = True
     return func
diff --git a/packages/cubejs-backend-native/src/python/entry.rs b/packages/cubejs-backend-native/src/python/entry.rs
index 5dc78b1ebd004..043c56dbe3319 100644
--- a/packages/cubejs-backend-native/src/python/entry.rs
+++ b/packages/cubejs-backend-native/src/python/entry.rs
@@ -73,7 +73,21 @@ fn python_load_model(mut cx: FunctionContext) -> JsResult<JsPromise> {
         let model_module = PyModule::from_code(py, &model_content, &model_file_name, "")?;
         let mut collected_functions = CLReprObject::new();
 
-        if model_module.hasattr("__execution_context_locals")? {
+        if model_module.hasattr("template")? {
+            let functions = model_module
+                .getattr("template")?
+                .getattr("functions")?
+                .downcast::<PyDict>()?;
+
+            for (local_key, local_value) in functions.iter() {
+                if local_value.is_instance_of::<PyFunction>() {
+                    let fun: Py<PyFunction> = local_value.downcast::<PyFunction>()?.into();
+                    collected_functions
+                        .insert(local_key.to_string(), CLRepr::PyExternalFunction(fun));
+                }
+            }
+            // TODO remove all other ways of defining functions
+        } else if model_module.hasattr("__execution_context_locals")? {
             let execution_context_locals = model_module
                 .getattr("__execution_context_locals")?
                 .downcast::<PyDict>()?;
diff --git a/packages/cubejs-backend-native/test/templates/jinja-instance.py b/packages/cubejs-backend-native/test/templates/jinja-instance.py
index 063da194b64e5..18303610770ef 100644
--- a/packages/cubejs-backend-native/test/templates/jinja-instance.py
+++ b/packages/cubejs-backend-native/test/templates/jinja-instance.py
@@ -6,8 +6,8 @@
 def arg_sum_integers(a, b):
   return a + b
 
-@template.function
-def arg_bool(a):
+@template.function("arg_bool")
+def ab(a):
   return a + 0
 
 @template.function