From e062c40679c8a82f771a6ea81854de0225e9fc53 Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 10 Dec 2024 14:35:41 -0500 Subject: [PATCH 1/6] Add syntax support for types on macro parameters. --- dbt_common/clients/jinja.py | 38 +++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/dbt_common/clients/jinja.py b/dbt_common/clients/jinja.py index 8cc1205..7bfc5e1 100644 --- a/dbt_common/clients/jinja.py +++ b/dbt_common/clients/jinja.py @@ -103,6 +103,44 @@ def parse_macro(self) -> jinja2.nodes.Macro: node.body = self.parse_statements(("name:endmacro",), drop_needle=True) return node + def parse_signature(self, node: Union[jinja2.nodes.Macro, jinja2.nodes.CallBlock]) -> None: + """Overrides the default jinja Parser.parse_signature method, modifying + the original implementation to allow macros to have typed parameters.""" + + # Jinja does not support extending its node types, such as Macro, so + # at least while typed macros are experimental, we will patch the + # information onto the existing types. + setattr(node, "arg_types", []) + setattr(node, "has_type_annotations", False) + + args = node.args = [] + defaults = node.defaults = [] + + self.stream.expect("lparen") + while self.stream.current.type != "rparen": + if args: + self.stream.expect("comma") + + arg = self.parse_assign_target(name_only=True) + arg.set_ctx("param") + + type_name: Optional[str] + if self.stream.skip_if("colon"): + node.has_type_annotations = True # type: ignore + type_name = self.stream.expect("name") + else: + type_name = "" + + node.arg_types.append(type_name) # type: ignore + + if self.stream.skip_if("assign"): + defaults.append(self.parse_expression()) + elif defaults: + self.fail("non-default argument follows default argument") + + args.append(arg) + self.stream.expect("rparen") + class MacroFuzzEnvironment(jinja2.sandbox.SandboxedEnvironment): def _parse( From 5579337daed88ecc5b4d8d101b1f4461e624073e Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 10 Dec 2024 14:36:40 -0500 Subject: [PATCH 2/6] Fixes for mypy --- dbt_common/clients/jinja.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dbt_common/clients/jinja.py b/dbt_common/clients/jinja.py index 7bfc5e1..ac3afa9 100644 --- a/dbt_common/clients/jinja.py +++ b/dbt_common/clients/jinja.py @@ -113,8 +113,8 @@ def parse_signature(self, node: Union[jinja2.nodes.Macro, jinja2.nodes.CallBlock setattr(node, "arg_types", []) setattr(node, "has_type_annotations", False) - args = node.args = [] - defaults = node.defaults = [] + args = node.args = [] # type: ignore + defaults = node.defaults = [] # type: ignore self.stream.expect("lparen") while self.stream.current.type != "rparen": From 8eeb25e927ea8bc958f3f98a5b82c6af136ad868 Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 10 Dec 2024 16:16:54 -0500 Subject: [PATCH 3/6] Add support for parameterized types. Add unit tests. --- dbt_common/clients/jinja.py | 30 +++++++++++++++++++++++++++++- tests/unit/test_jinja.py | 36 +++++++++++++++++++++++++++++++++++- 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/dbt_common/clients/jinja.py b/dbt_common/clients/jinja.py index ac3afa9..4bba253 100644 --- a/dbt_common/clients/jinja.py +++ b/dbt_common/clients/jinja.py @@ -1,4 +1,5 @@ import codecs +import dataclasses import linecache import os import tempfile @@ -90,6 +91,12 @@ def _linecache_inject(source: str, write: bool) -> str: return filename +@dataclasses.dataclass +class MacroType: + name: str + type_params: List["MacroType"] = dataclasses.field(default_factory=list) + + class MacroFuzzParser(jinja2.parser.Parser): def parse_macro(self) -> jinja2.nodes.Macro: node = jinja2.nodes.Macro(lineno=next(self.stream).lineno) @@ -127,7 +134,7 @@ def parse_signature(self, node: Union[jinja2.nodes.Macro, jinja2.nodes.CallBlock type_name: Optional[str] if self.stream.skip_if("colon"): node.has_type_annotations = True # type: ignore - type_name = self.stream.expect("name") + type_name = self.parse_type_name() else: type_name = "" @@ -141,6 +148,27 @@ def parse_signature(self, node: Union[jinja2.nodes.Macro, jinja2.nodes.CallBlock args.append(arg) self.stream.expect("rparen") + def parse_type_name(self) -> MacroType: + # NOTE: Types syntax is validated here, but not whether type names + # are valid or have correct parameters. + + # A type name should consist of a name (i.e. 'Dict')... + type_name = self.stream.expect("name").value + type = MacroType(type_name) + + # ..and an optional comma-delimited list of type parameters + # as in the type declaration 'Dict[str, str]' + if self.stream.skip_if("lbracket"): + while self.stream.current.type != "rbracket": + if type.type_params: + self.stream.expect("comma") + param_type = self.parse_type_name() + type.type_params.append(param_type) + + self.stream.expect("rbracket") + + return type + class MacroFuzzEnvironment(jinja2.sandbox.SandboxedEnvironment): def _parse( diff --git a/tests/unit/test_jinja.py b/tests/unit/test_jinja.py index e839296..259f614 100644 --- a/tests/unit/test_jinja.py +++ b/tests/unit/test_jinja.py @@ -1,7 +1,8 @@ +import jinja2 import unittest from dbt_common.clients._jinja_blocks import BlockTag -from dbt_common.clients.jinja import extract_toplevel_blocks, get_template, render_template +from dbt_common.clients.jinja import extract_toplevel_blocks, get_template, render_template, MacroFuzzParser, MacroType from dbt_common.exceptions import CompilationError @@ -524,3 +525,36 @@ def test_if_list_filter(): template = get_template(jinja_string, ctx) rendered = render_template(template, ctx) assert "Did not find a list" in rendered + + +def test_macro_parser_parses_simple_types() -> None: + macro_txt = """ + {% macro test_macro(param1: str, param2: int, param3: bool, param4: float, param5: Any) %} + {% endmacro %} + """ + + env = jinja2.Environment() + parser = MacroFuzzParser(env, macro_txt) + result = parser.parse() + arg_types = result.body[1].arg_types + assert arg_types[0] == MacroType("str") + assert arg_types[1] == MacroType("int") + assert arg_types[2] == MacroType("bool") + assert arg_types[3] == MacroType("float") + assert arg_types[4] == MacroType("Any") + + +def test_macro_parser_parses_complex_types() -> None: + macro_txt = """ + {% macro test_macro(param1: List[str], param2: Dict[ int,str ], param3: Optional[List[str]], param4: Dict[str, Dict[bool, Any]]) %} + {% endmacro %} + """ + + env = jinja2.Environment() + parser = MacroFuzzParser(env, macro_txt) + result = parser.parse() + arg_types = result.body[1].arg_types + assert arg_types[0] == MacroType("List", [MacroType("str")]) + assert arg_types[1] == MacroType("Dict", [MacroType("int"), MacroType("str")]) + assert arg_types[2] == MacroType("Optional", [MacroType("List", [MacroType("str")])]) + assert arg_types[3] == MacroType("Dict", [MacroType("str"), MacroType("Dict", [MacroType("bool"), MacroType("Any")])]) From 30c5b97842709db2c50bb3d4ab029617c9e675f2 Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 10 Dec 2024 16:23:13 -0500 Subject: [PATCH 4/6] Fix formatting. --- tests/unit/test_jinja.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_jinja.py b/tests/unit/test_jinja.py index 259f614..3a35abf 100644 --- a/tests/unit/test_jinja.py +++ b/tests/unit/test_jinja.py @@ -2,7 +2,13 @@ import unittest from dbt_common.clients._jinja_blocks import BlockTag -from dbt_common.clients.jinja import extract_toplevel_blocks, get_template, render_template, MacroFuzzParser, MacroType +from dbt_common.clients.jinja import ( + extract_toplevel_blocks, + get_template, + render_template, + MacroFuzzParser, + MacroType, +) from dbt_common.exceptions import CompilationError @@ -506,7 +512,7 @@ def test_if_endfor_newlines(self) -> None: """ -def test_if_list_filter(): +def test_if_list_filter() -> None: jinja_string = """ {%- if my_var | is_list -%} Found a list @@ -557,4 +563,6 @@ def test_macro_parser_parses_complex_types() -> None: assert arg_types[0] == MacroType("List", [MacroType("str")]) assert arg_types[1] == MacroType("Dict", [MacroType("int"), MacroType("str")]) assert arg_types[2] == MacroType("Optional", [MacroType("List", [MacroType("str")])]) - assert arg_types[3] == MacroType("Dict", [MacroType("str"), MacroType("Dict", [MacroType("bool"), MacroType("Any")])]) + assert arg_types[3] == MacroType( + "Dict", [MacroType("str"), MacroType("Dict", [MacroType("bool"), MacroType("Any")])] + ) From f72d183dcf79ba6df0831bf08c7aac4e8bc7261d Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 10 Dec 2024 16:24:29 -0500 Subject: [PATCH 5/6] Add changelog entry --- .changes/unreleased/Features-20241210-144247.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .changes/unreleased/Features-20241210-144247.yaml diff --git a/.changes/unreleased/Features-20241210-144247.yaml b/.changes/unreleased/Features-20241210-144247.yaml new file mode 100644 index 0000000..439afb4 --- /dev/null +++ b/.changes/unreleased/Features-20241210-144247.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Add syntax support for types on macro parameters. +time: 2024-12-10T14:42:47.253157-05:00 +custom: + Author: peterallenwebb + Issue: "229" From 525700c26a4e4ff05feee80c0b7d91ab819957db Mon Sep 17 00:00:00 2001 From: Peter Allen Webb Date: Tue, 10 Dec 2024 16:27:29 -0500 Subject: [PATCH 6/6] Fix type annotations --- tests/unit/test_jinja.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_jinja.py b/tests/unit/test_jinja.py index 3a35abf..2247374 100644 --- a/tests/unit/test_jinja.py +++ b/tests/unit/test_jinja.py @@ -1,6 +1,8 @@ import jinja2 import unittest +from typing import Any, Dict + from dbt_common.clients._jinja_blocks import BlockTag from dbt_common.clients.jinja import ( extract_toplevel_blocks, @@ -521,7 +523,7 @@ def test_if_list_filter() -> None: {%- endif -%} """ # Check with list variable - ctx = {"my_var": ["one", "two"]} + ctx: Dict[str, Any] = {"my_var": ["one", "two"]} template = get_template(jinja_string, ctx) rendered = render_template(template, ctx) assert "Found a list" in rendered