Allow non-string types in variable templates and add magic omit value

This commit is contained in:
Tulir Asokan 2023-10-05 21:54:11 +03:00
parent 299cfdff46
commit a3baf06ca0
3 changed files with 25 additions and 12 deletions

View file

@ -16,7 +16,8 @@
from typing import List, Union, Dict, Any from typing import List, Union, Dict, Any
import re 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.util.config import BaseProxyConfig, ConfigUpdateHelper
from mautrix.types import EventType from mautrix.types import EventType
@ -92,17 +93,17 @@ class Config(BaseProxyConfig):
@staticmethod @staticmethod
def _parse_variables(data: Dict[str, Any]) -> Dict[str, Any]: 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("{{") if isinstance(var_tpl, str) and var_tpl.startswith("{{")
else var_tpl) else var_tpl)
for name, var_tpl in data.get("variables", {}).items()} for name, var_tpl in data.get("variables", {}).items()}
@staticmethod @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: if not content:
return {} return {}
elif isinstance(content, str): elif isinstance(content, str):
return JinjaTemplate(content) return JinjaStringTemplate(content)
return content return content
@staticmethod @staticmethod

View file

@ -16,13 +16,12 @@
from typing import Optional, Match, Dict, List, Set, Union, Pattern, Any from typing import Optional, Match, Dict, List, Set, Union, Pattern, Any
from attr import dataclass from attr import dataclass
from jinja2 import Template as JinjaTemplate
from mautrix.types import RoomID, EventType from mautrix.types import RoomID, EventType
from maubot import MessageEvent from maubot import MessageEvent
from .template import Template from .template import Template, OmitValue
from .simplepattern import SimplePattern from .simplepattern import SimplePattern
RPattern = Union[Pattern, SimplePattern] RPattern = Union[Pattern, SimplePattern]

View file

@ -20,7 +20,8 @@ import copy
import re import re
from attr import dataclass 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 from mautrix.types import EventType, Event
@ -30,6 +31,11 @@ class Key(str):
variable_regex = re.compile(r"\$\${([0-9A-Za-z-_]+)}") variable_regex = re.compile(r"\$\${([0-9A-Za-z-_]+)}")
OmitValue = object()
global_vars = {
"omit": OmitValue,
}
Index = Union[str, int, Key] Index = Union[str, int, Key]
@ -38,7 +44,7 @@ Index = Union[str, int, Key]
class Template: class Template:
type: EventType type: EventType
variables: Dict[str, Any] variables: Dict[str, Any]
content: Union[Dict[str, Any], JinjaTemplate] content: Union[Dict[str, Any], JinjaStringTemplate]
_variable_locations: List[Tuple[Index, ...]] = None _variable_locations: List[Tuple[Index, ...]] = None
@ -78,11 +84,14 @@ class Template:
) -> Dict[str, Any]: ) -> Dict[str, Any]:
variables = extra_vars variables = extra_vars
for name, template in chain(rule_vars.items(), self.variables.items()): for name, template in chain(rule_vars.items(), self.variables.items()):
if isinstance(template, JinjaTemplate): if isinstance(template, JinjaNativeTemplate):
variables[name] = template.render(event=evt, variables=variables) 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: else:
variables[name] = template variables[name] = template
if isinstance(self.content, JinjaTemplate): if isinstance(self.content, JinjaStringTemplate):
raw_json = self.content.render(event=evt, **variables) raw_json = self.content.render(event=evt, **variables)
return json.loads(raw_json) return json.loads(raw_json)
content = copy.deepcopy(self.content) content = copy.deepcopy(self.content)
@ -93,5 +102,9 @@ class Template:
key = str(key) key = str(key)
data[self._replace_variables(key, variables)] = data.pop(key) data[self._replace_variables(key, variables)] = data.pop(key)
else: 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 return content