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
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

View file

@ -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]

View file

@ -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