From 8eeb25e927ea8bc958f3f98a5b82c6af136ad868 Mon Sep 17 00:00:00 2001
From: Peter Allen Webb <peter.webb@dbtlabs.com>
Date: Tue, 10 Dec 2024 16:16:54 -0500
Subject: [PATCH] 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")])])