diff --git a/tealish/nodes.py b/tealish/nodes.py index 70f236f..22de921 100644 --- a/tealish/nodes.py +++ b/tealish/nodes.py @@ -231,6 +231,8 @@ def consume(cls, compiler: "TealishCompiler", parent: Node) -> "Statement": return Switch.consume(compiler, parent) elif line.startswith("func "): return Func.consume(compiler, parent) + elif line.startswith("@"): + return DecoratedFunc.consume(compiler, parent) elif line.startswith("if "): return IfStatement.consume(compiler, parent) elif line.startswith("while "): @@ -1478,6 +1480,8 @@ def __init__( except KeyError as e: raise ParseError(str(e) + f" Line {self.line_no}") self.vars: Dict[str, Var] = {} + self.decorators = [] + self.attributes = {} @classmethod def consume(cls, compiler: "TealishCompiler", parent: Optional[Node]) -> "Func": @@ -1579,6 +1583,76 @@ def _tealish(self) -> str: return output + "\n" +class Decorator(Node): + pattern = r"@(?P[a-z][a-zA-Z_0-9]*)\((?P.*)\)$" + name: str + params: str + + def write_teal(self, writer: "TealWriter") -> None: + writer.write(self, f"// {self.line}") + + def _tealish(self) -> str: + output = f"@{self.name}({self.params})\n" + return output + + +class DecoratedFunc(InlineStatement): + possible_child_nodes = [Decorator, Comment, Func] + pattern = r"" + + def __init__( + self, + line: str, + parent: Optional[Node] = None, + compiler: Optional["TealishCompiler"] = None, + ) -> None: + super().__init__(line, parent, compiler) + self.func = None + self.decorators = [] + + def add_decorator(self, node) -> None: + self.decorators.append(node) + self.add_child(node) + + def set_func(self, node) -> None: + self.func = node + self.func.decorators = self.decorators + for decorator in self.decorators: + self.func.attributes[decorator.name] = {} + if m := re.match(r"(?P.*)=(?P.*)", decorator.params): + self.func.attributes[decorator.name] = { + m.groupdict()["key"]: m.groupdict()["value"] + } + self.add_child(node) + + @classmethod + def consume( + cls, compiler: "TealishCompiler", parent: Optional[Node] + ) -> "DecoratedFunc": + decorated_func = DecoratedFunc("", parent, compiler=compiler) + while True: + if compiler.peek().startswith("func "): + decorated_func.set_func(Func.consume(compiler, decorated_func)) + break + elif compiler.peek().startswith("@"): + decorated_func.add_decorator(Decorator.consume(compiler, parent)) + return decorated_func + + def process(self) -> None: + for node in self.child_nodes: + node.process() + + def write_teal(self, writer: "TealWriter") -> None: + writer.write(self, self.func) + + def _tealish(self) -> str: + output = "" + for n in self.decorators: + output += n.tealish() + output += self.func.tealish() + return output + + class StructFieldDefinition(InlineStatement): pattern = ( r"(?P[a-z][A-Z-a-z0-9_]*): "