From a3baf06ca0c700a07af0d9d0d4ea6c436debe105 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Thu, 5 Oct 2023 21:54:11 +0300
Subject: [PATCH] Allow non-string types in variable templates and add magic
 omit value

---
 reactbot/config.py   |  9 +++++----
 reactbot/rule.py     |  3 +--
 reactbot/template.py | 25 +++++++++++++++++++------
 3 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/reactbot/config.py b/reactbot/config.py
index efbf9ae..83e2283 100644
--- a/reactbot/config.py
+++ b/reactbot/config.py
@@ -16,7 +16,8 @@
 from typing import List, Union, Dict, Any
 import re
 
-from jinja2 import Template as JinjaTemplate
+from jinja2 import Template as JinjaStringTemplate
+from jinja2.nativetypes import NativeTemplate as JinjaNativeTemplate
 
 from mautrix.util.config import BaseProxyConfig, ConfigUpdateHelper
 from mautrix.types import EventType
@@ -92,17 +93,17 @@ class Config(BaseProxyConfig):
 
     @staticmethod
     def _parse_variables(data: Dict[str, Any]) -> Dict[str, Any]:
-        return {name: (JinjaTemplate(var_tpl)
+        return {name: (JinjaNativeTemplate(var_tpl)
                        if isinstance(var_tpl, str) and var_tpl.startswith("{{")
                        else var_tpl)
                 for name, var_tpl in data.get("variables", {}).items()}
 
     @staticmethod
-    def _parse_content(content: Union[Dict[str, Any], str]) -> Union[Dict[str, Any], JinjaTemplate]:
+    def _parse_content(content: Union[Dict[str, Any], str]) -> Union[Dict[str, Any], JinjaStringTemplate]:
         if not content:
             return {}
         elif isinstance(content, str):
-            return JinjaTemplate(content)
+            return JinjaStringTemplate(content)
         return content
 
     @staticmethod
diff --git a/reactbot/rule.py b/reactbot/rule.py
index b97b350..51a8bee 100644
--- a/reactbot/rule.py
+++ b/reactbot/rule.py
@@ -16,13 +16,12 @@
 from typing import Optional, Match, Dict, List, Set, Union, Pattern, Any
 
 from attr import dataclass
-from jinja2 import Template as JinjaTemplate
 
 from mautrix.types import RoomID, EventType
 
 from maubot import MessageEvent
 
-from .template import Template
+from .template import Template, OmitValue
 from .simplepattern import SimplePattern
 
 RPattern = Union[Pattern, SimplePattern]
diff --git a/reactbot/template.py b/reactbot/template.py
index 7643ca4..f150a59 100644
--- a/reactbot/template.py
+++ b/reactbot/template.py
@@ -20,7 +20,8 @@ import copy
 import re
 
 from attr import dataclass
-from jinja2 import Template as JinjaTemplate
+from jinja2 import Template as JinjaStringTemplate
+from jinja2.nativetypes import Template as JinjaNativeTemplate
 
 from mautrix.types import EventType, Event
 
@@ -30,6 +31,11 @@ class Key(str):
 
 
 variable_regex = re.compile(r"\$\${([0-9A-Za-z-_]+)}")
+OmitValue = object()
+
+global_vars = {
+    "omit": OmitValue,
+}
 
 Index = Union[str, int, Key]
 
@@ -38,7 +44,7 @@ Index = Union[str, int, Key]
 class Template:
     type: EventType
     variables: Dict[str, Any]
-    content: Union[Dict[str, Any], JinjaTemplate]
+    content: Union[Dict[str, Any], JinjaStringTemplate]
 
     _variable_locations: List[Tuple[Index, ...]] = None
 
@@ -78,11 +84,14 @@ class Template:
                 ) -> Dict[str, Any]:
         variables = extra_vars
         for name, template in chain(rule_vars.items(), self.variables.items()):
-            if isinstance(template, JinjaTemplate):
-                variables[name] = template.render(event=evt, variables=variables)
+            if isinstance(template, JinjaNativeTemplate):
+                rendered_var = template.render(event=evt, variables=variables, **global_vars)
+                if not isinstance(rendered_var, (str, int, list, tuple, dict, bool)) and rendered_var is not None and rendered_var is not OmitValue:
+                    rendered_var = str(rendered_var)
+                variables[name] = rendered_var
             else:
                 variables[name] = template
-        if isinstance(self.content, JinjaTemplate):
+        if isinstance(self.content, JinjaStringTemplate):
             raw_json = self.content.render(event=evt, **variables)
             return json.loads(raw_json)
         content = copy.deepcopy(self.content)
@@ -93,5 +102,9 @@ class Template:
                 key = str(key)
                 data[self._replace_variables(key, variables)] = data.pop(key)
             else:
-                data[key] = self._replace_variables(data[key], variables)
+                replaced_data = self._replace_variables(data[key], variables)
+                if replaced_data is OmitValue:
+                    del data[key]
+                else:
+                    data[key] = replaced_data
         return content