Repository: TobikoData/sqlmesh
Pull request #2261 (2024-03-13 17:22:47)

Toby/revert types


Patch:
diff --git a/docs/concepts/macros/sqlmesh_macros.md b/docs/concepts/macros/sqlmesh_macros.md
index bae6415c..22ea3d73 100644
--- a/docs/concepts/macros/sqlmesh_macros.md
+++ b/docs/concepts/macros/sqlmesh_macros.md
@@ -1044,116 +1044,6 @@ def some_macro(evaluator):
     ...
 ```
 
-## Typed Macros
-
-Typed macros in SQLMesh bring the power of type hints from Python, enhancing readability, maintainability, and usability of your SQL macros. These macros enable developers to specify expected types for arguments, making the macros more intuitive and less error-prone.
-
-### Benefits of Typed Macros
-
-1. **Improved Readability**: By specifying types, the intent of the macro is clearer to other developers or future you.
-2. **Reduced Boilerplate**: No need for manual type conversion within the macro function, allowing you to focus on the core logic.
-3. **Enhanced Autocompletion**: IDEs can provide better autocompletion and documentation based on the specified types.
-
-### Defining a Typed Macro
-
-Typed macros in SQLMesh use Python's type hints. Here's a simple example of a typed macro that repeats a string a given number of times:
-
-```python linenums="1"
-from sqlmesh import macro
-
-@macro()
-def repeat_string(evaluator, text: str, count: int) -> str:
-    return text * count
-```
-
-Usage in SQLMesh:
-
-```sql linenums="1"
-SELECT
-  @repeat_string('SQLMesh ', 3) as repeated_string
-FROM some_table;
-```
-
-This macro takes two arguments: `text` of type `str` and `count` of type `int`, and it returns a string. Without type hints, the inputs to the macro would have been two `exp.Literal` objects you would have had to convert to strings and integers manually.
-
-### Supported Types
-
-SQLMesh supports common Python types for typed macros including:
-
-- `str`
-- `int`
-- `float`
-- `bool`
-- `List[T]` - where `T` is any supported type including sqlglot expressions
-- `Tuple[T]` - where `T` is any supported type including sqlglot expressions
-- `Union[T1, T2, ...]` - where `T1`, `T2`, etc. are any supported types including sqlglot expressions
-
-We also support SQLGlot expressions as type hints, allowing you to ensure inputs are coerced to the desired SQL AST node your intending on working with. Some useful examples include:
-
-- `exp.Table`
-- `exp.Column`
-- `exp.Literal`
-- `exp.Identifier`
-
-While these might be obvious examples, you can effectively coerce an input into _any_ SQLGlot expression type, which can be useful for more complex macros. When coercing to more complex types, you will almost certainly need to pass a string literal since expression to expression coercion is limited. When a string literal is passed to a macro that hints at a SQLGlot expression, the string will be parsed using SQLGlot and coerced to the correct type. Failure to coerce to the correct type will result in the original expression being passed to the macro and a warning being logged for the user to address as-needed.
-
-```python linenums="1"
-@macro()
-def stamped(evaluator, query: exp.Select) -> exp.Subquery:
-    return query.select(exp.Literal.string(str(datetime.now())).as_("stamp")).subquery()
-
-# Coercing to a complex node like `exp.Select` works as expected given a string literal input
-# SELECT * FROM @stamped('SELECT a, b, c')
-```
-
-When coercion fails, there will always be a warning logged but we will not crash. We believe the macro should be flexible by default, meaning the default behavior is preserved if we cannot coerce. Give that, the user can express whatever level of additional checks they want. For example, if you would like to raise an error when the coercion fails, you can use an `assert` statement. For example:
-
-```python linenums="1"
-@macro()
-def my_macro(evaluator, table: exp.Table) -> exp.Column:
-    assert isinstance(table, exp.Table)
-    table.set("catalog", "dev")
-    return table
-
-# Works
-# SELECT * FROM @my_macro('some.table')
-# SELECT * FROM @my_macro(some.table)
-
-# Raises an error thanks to the users inclusion of the assert, otherwise would pass through the string literal and log a warning
-# SELECT * FROM @my_macro('SELECT 1 + 1')
-```
-
-In using assert this way, you still get the benefits of reducing/removing the boilerplate needed to coerce types; but you **also** get guarantees about the type of the input. This is a useful pattern and is user-defined, so you can use it as you see fit. It ultimately allows you to keep the macro definition clean and focused on the core business logic.
-
-### Advanced Typed Macros
-
-You can create more complex macros using advanced Python features like generics. For example, a macro that accepts a list of integers and returns their sum:
-
-```python linenums="1"
-from typing import List
-from sqlmesh import macro
-
-@macro()
-def sum_integers(evaluator, numbers: List[int]) -> int:
-    return sum(numbers)
-```
-
-Usage in SQLMesh:
-
-```sql linenums="1"
-SELECT
-  @sum_integers([1, 2, 3, 4, 5]) as total
-FROM some_table;
-```
-
-Generics can be nested and are resolved recursively allowing for fairly robust type hinting.
-
-See examples of the coercion function in action in the test suite [here](../../../tests/core/test_macros.py).
-
-### Conclusion
-
-Typed macros in SQLMesh not only enhance the development experience by making macros more readable and easier to use but also contribute to more robust and maintainable code. By leveraging Python's type hinting system, developers can create powerful and intuitive macros for their SQL queries, further bridging the gap between SQL and Python.
-
 ## Mixing macro systems
 
 SQLMesh supports both SQLMesh and [Jinja](./jinja_macros.md) macro systems. We strongly recommend using only one system in a model - if both are present, they may fail or behave in unintuitive ways.
diff --git a/examples/sushi/macros/macros.py b/examples/sushi/macros/macros.py
index f4c2b6a8..abf0ebb9 100644
--- a/examples/sushi/macros/macros.py
+++ b/examples/sushi/macros/macros.py
@@ -1,9 +1,8 @@
 from macros.utils import between  # type: ignore
-from sqlglot import exp
 
 from sqlmesh import macro
 
 
 @macro()
-def incremental_by_ds(evaluator, column: exp.Column):
+def incremental_by_ds(evaluator, column):
     return between(evaluator, column, evaluator.locals["start_date"], evaluator.locals["end_date"])
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index ddc3a7f2..6424400b 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -1,15 +1,13 @@
 from __future__ import annotations
 
-import inspect
-import logging
 import typing as t
 from enum import Enum
-from functools import reduce, wraps
+from functools import reduce
 from string import Template
 
 import sqlglot
 from jinja2 import Environment
-from sqlglot import Generator, exp, parse_one
+from sqlglot import Generator, exp
 from sqlglot.executor.env import ENV
 from sqlglot.executor.python import Python
 from sqlglot.helper import csv, ensure_collection
@@ -17,7 +15,6 @@ from sqlglot.schema import MappingSchema
 
 from sqlmesh.core.dialect import (
     SQLMESH_MACRO_PREFIX,
-    Dialect,
     MacroDef,
     MacroFunc,
     MacroSQL,
@@ -36,8 +33,6 @@ if t.TYPE_CHECKING:
     from sqlmesh.core.engine_adapter import EngineAdapter
     from sqlmesh.core.snapshot import Snapshot
 
-logger = logging.getLogger(__name__)
-
 
 class RuntimeStage(Enum):
     LOADING = "loading"
@@ -354,75 +349,6 @@ class MacroEvaluator:
             )
         return self.locals["engine_adapter"]
 
-    def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.Any:
-        """Coerces the given expression to the specified type on a best-effort basis."""
-        base_err_msg = f"Failed to coerce expression '{expr}' to type '{typ}'."
-        try:
-            if typ is None or typ is t.Any:
-                return expr
-            base = t.get_origin(typ) or typ
-            # We need to handle t.Union first since we cannot use isinstance with it
-            if base is t.Union:
-                for branch in t.get_args(typ):
-                    try:
-                        return self._coerce(expr, branch, True)
-                    except Exception:
-                        pass
-                raise SQLMeshError(base_err_msg)
-            if isinstance(expr, base):
-                return expr
-            if issubclass(base, exp.Expression):
-                d = Dialect.get_or_raise(self.dialect)
-                into = base if base in d.parser().EXPRESSION_PARSERS else None
-                if into is None:
-                    if isinstance(expr, exp.Literal):
-                        coerced = parse_one(expr.this)
-                    else:
-                        raise SQLMeshError(
-                            f"{base_err_msg} Coercion to {base} requires a literal expression."
-                        )
-                else:
-                    coerced = parse_one(
-                        expr.this if isinstance(expr, exp.Literal) else expr.sql(), into=into
-                    )
-                if isinstance(coerced, base):
-                    return coerced
-                raise SQLMeshError(base_err_msg)
-
-            if base in (int, float, str) and isinstance(expr, exp.Literal):
-                return base(expr.this)
-            if base is bool and isinstance(expr, exp.Boolean):
-                return expr.this
-            if base is str and isinstance(expr, exp.Expression):
-                return expr.sql(self.dialect)
-            if base is tuple and isinstance(expr, (exp.Tuple, exp.Array)):
-                generic = t.get_args(typ)
-                if not generic:
-                    return tuple(expr.expressions)
-                if generic[-1] is ...:
-                    return tuple(self._coerce(expr, generic[0]) for expr in expr.expressions)
-                elif len(generic) == len(expr.expressions):
-                    return tuple(
-                        self._coerce(expr, generic[i]) for i, expr in enumerate(expr.expressions)
-                    )
-                raise SQLMeshError(f"{base_err_msg} Expected {len(generic)} items.")
-            if base is list and isinstance(expr, (exp.Array, exp.Tuple)):
-                generic = t.get_args(typ)
-                if not generic:
-                    return expr.expressions
-                return [self._coerce(expr, generic[0]) for expr in expr.expressions]
-            raise SQLMeshError(base_err_msg)
-        except Exception:
-            if strict:
-                raise
-            logger.warning(
-                "Coercion of expression '%s' to type '%s' failed. Using non coerced expression.",
-                expr,
-                typ,
-                exc_info=True,
-            )
-            return expr
-
 
 class macro(registry_decorator):
     """Specifies a function is a macro and registers it the global MACROS registry.
@@ -446,31 +372,7 @@ class macro(registry_decorator):
     def __call__(
         self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
     ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
-        @wraps(func)
-        def _typed_func(
-            evaluator: MacroEvaluator, *args_: t.Any, **kwargs_: t.Any
-        ) -> DECORATOR_RETURN_TYPE:
-            spec = inspect.getfullargspec(func)
-            annotations = t.get_type_hints(func)
-            kwargs = inspect.getcallargs(func, evaluator, *args_, **kwargs_)
-            for param, value in kwargs.items():
-                coercible_type = annotations.get(param)
-                if not coercible_type:
-                    continue
-                kwargs[param] = evaluator._coerce(value, coercible_type)
-            args = [kwargs.pop(k) for k in spec.args if k in kwargs]
-            if spec.varargs:
-                args.extend(kwargs.pop(spec.varargs, []))
-            return func(*args, **kwargs)
-
-        try:
-            annotated = any(t.get_type_hints(func).keys() - {"return"})
-        except TypeError:
-            annotated = False
-
-        wrapper = super().__call__(
-            func if not annotated else t.cast(t.Callable[..., DECORATOR_RETURN_TYPE], _typed_func)
-        )
+        wrapper = super().__call__(func)
 
         # This is useful to identify macros at runtime
         setattr(wrapper, "__sqlmesh_macro__", True)
diff --git a/sqlmesh/utils/metaprogramming.py b/sqlmesh/utils/metaprogramming.py
index 6d24399f..16a0619e 100644
--- a/sqlmesh/utils/metaprogramming.py
+++ b/sqlmesh/utils/metaprogramming.py
@@ -215,6 +215,8 @@ def normalize_source(obj: t.Any) -> str:
             # remove function return type annotation
             if isinstance(node, ast.FunctionDef):
                 node.returns = None
+        elif isinstance(node, ast.arg):
+            node.annotation = None
 
     return to_source(root_node).strip()
 
@@ -276,8 +278,7 @@ def build_env(
     if name not in env:
         # We only need to add the undecorated code of @macro() functions in env, which
         # is accessible through the `__wrapped__` attribute added by functools.wraps
-        # We account for the case where the function is wrapped multiple times too
-        env[name] = inspect.unwrap(obj) if hasattr(obj, "__sqlmesh_macro__") else obj
+        env[name] = obj.__wrapped__ if getattr(obj, "__sqlmesh_macro__", None) else obj
 
         if obj_module and _is_relative_to(obj_module.__file__, path):
             walk(env[name])


Test patch:
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index ee6c2d9e..4302ef92 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -1,7 +1,5 @@
-import functools
 import sys
 import typing as t
-from pathlib import Path
 from unittest import mock
 
 import pytest
@@ -10,12 +8,7 @@ from sqlglot import exp, parse_one
 from sqlmesh.core.dialect import StagedFilePath
 from sqlmesh.core.macros import MacroEvaluator, macro
 from sqlmesh.utils.errors import SQLMeshError
-from sqlmesh.utils.metaprogramming import (
-    Executable,
-    ExecutableKind,
-    build_env,
-    serialize_env,
-)
+from sqlmesh.utils.metaprogramming import Executable, ExecutableKind
 
 
 @pytest.fixture
@@ -34,45 +27,6 @@ def macro_evaluator() -> MacroEvaluator:
     def noop(evaluator: MacroEvaluator):
         return None
 
-    @macro()
-    def bitshift_square(evaluator: MacroEvaluator, x: int, y: int) -> int:
-        return (x >> y) ** 2
-
-    @macro()
-    def prefix_db(evaluator: MacroEvaluator, table: exp.Table, prefix: str) -> exp.Table:
-        table.set("db", prefix + table.db)
-        return table
-
-    @macro()
-    def repeated(evaluator: MacroEvaluator, expr: str, times: int = 2, multi: bool = False):
-        if multi is True:
-            return (expr,) * times
-        return expr * times
-
-    @macro()
-    def split(evaluator: MacroEvaluator, string: str, sep: str = ","):
-        return string.split(sep)
-
-    @macro()
-    def cte_tag_name(evaluator: MacroEvaluator, with_: exp.Select):
-        for cte in with_.find_all(exp.CTE):
-            name = cte.alias_or_name
-            for query in cte.find_all(exp.Select):
-                query.select(exp.Literal.string(name).as_("source"), copy=False)
-        return with_
-
-    @macro()
-    def suffix_idents(evaluator: MacroEvaluator, items: t.List[str], suffix: str):
-        return [item + suffix for item in items]
-
-    @macro()
-    def suffix_idents_2(evaluator: MacroEvaluator, items: t.Tuple[str, ...], suffix: str):
-        return [item + suffix for item in items]
-
-    @macro()
-    def stamped(evaluator, query: exp.Select) -> exp.Subquery:
-        return query.select(exp.Literal.string("2024-01-01").as_("stamp")).subquery()
-
     return MacroEvaluator(
         "hive",
         {"test": Executable(name="test", payload=f"def test(_):\n    return 'test'")},
@@ -87,15 +41,12 @@ def test_external_macro() -> None:
     def foo(evaluator: MacroEvaluator) -> str:
         return "foo"
 
-    @functools.wraps(foo)
-    def wrapper(*args, **kwargs):
-        return foo(*args, **kwargs)
-
     # Mimic a SQLMesh macro definition by setting the func's metadata appropriately
-    setattr(wrapper, "__sqlmesh_macro__", True)
+    setattr(foo, "__wrapped__", foo)
+    setattr(foo, "__sqlmesh_macro__", True)
 
     sys.modules["pkg"] = mock.Mock()
-    with mock.patch("pkg.foo", wrapper):
+    with mock.patch("pkg.foo", foo):
         evaluator = MacroEvaluator(
             python_env={
                 "foo": Executable(
@@ -111,31 +62,6 @@ def test_external_macro() -> None:
         assert "@FOO" in evaluator.macros
         assert evaluator.macros["@FOO"](evaluator) == "foo"
 
-    # Assert that serializing a wrapped external macro traverses the chain of wrapped functions
-    # even through other modules
-    fakemod = mock.Mock()
-    fakemod.__name__ = "other_pkg"
-    fakemod.__file__ = "some/path.py"
-    sys.modules["other_pkg"] = fakemod
-
-    @functools.wraps(wrapper)
-    def nested_wrapper(*args, **kwargs):
-        return wrapper(*args, **kwargs)
-
-    setattr(nested_wrapper, "__sqlmesh_macro__", True)
-    setattr(nested_wrapper, "__module__", "other_pkg")
-    with mock.patch("other_pkg.bar", nested_wrapper):
-        env: t.Dict[str, t.Any] = {}
-        build_env(nested_wrapper, env=env, name="bar", path=Path("tests/core/test_macros.py"))
-        assert serialize_env(env, path=Path("tests/core")) == {
-            "bar": Executable(
-                payload="def foo(evaluator: MacroEvaluator):\n    return 'foo'",
-                name="foo",
-                path="test_macros.py",
-                alias="bar",
-            )
-        }
-
 
 def test_macro_var(macro_evaluator):
     expression = parse_one("@x")
@@ -366,56 +292,6 @@ def test_ast_correctness(macro_evaluator):
             "SELECT d",
             {},
         ),
-        (
-            """@BITSHIFT_SQUARE(50, '3')""",
-            "36",
-            {},
-        ),
-        (
-            """@PREFIX_DB(my.schema.table, 'dev_')""",
-            "my.dev_schema.table",
-            {},
-        ),
-        (
-            """select @REPEATED(test, 3)""",
-            "SELECT testtesttest",
-            {},
-        ),
-        (
-            """select @REPEATED(test, 3, true)""",
-            "SELECT test, test, test",
-            {},
-        ),
-        (
-            """select @SPLIT('a,b,c')""",
-            "SELECT a, b, c",
-            {},
-        ),
-        (
-            """@CTE_TAG_NAME(WITH step1 AS (SELECT 1) SELECT * FROM step1)""",
-            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
-            {},
-        ),
-        (
-            """@CTE_TAG_NAME('WITH step1 AS (SELECT 1) SELECT * FROM step1')""",
-            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
-            {},
-        ),
-        (
-            """SELECT @SUFFIX_IDENTS(['a', 'b', 'c'], 'z')""",
-            "SELECT az, bz, cz",
-            {},
-        ),
-        (
-            """SELECT @SUFFIX_IDENTS_2(['a', 'b', 'c'], 'z')""",
-            "SELECT az, bz, cz",
-            {},
-        ),
-        (
-            """select * FROM @STAMPED('select a, b, c')""",
-            "SELECT * FROM (SELECT a, b, c, '2024-01-01' AS stamp)",
-            {},
-        ),
         (
             """@OR(NULL, TRUE)""",
             "TRUE",
@@ -443,84 +319,10 @@ def test_ast_correctness(macro_evaluator):
         ),
     ],
 )
-def test_macro_functions(macro_evaluator: MacroEvaluator, assert_exp_eq, sql, expected, args):
+def test_macro_functions(macro_evaluator, assert_exp_eq, sql, expected, args):
     macro_evaluator.locals = args or {}
     assert_exp_eq(macro_evaluator.transform(parse_one(sql)), expected)
 
 
-def test_macro_returns_none(macro_evaluator: MacroEvaluator):
+def test_macro_returns_none(macro_evaluator):
     assert macro_evaluator.transform(parse_one("@NOOP()")) is None
-
-
-def test_macro_coercion(macro_evaluator: MacroEvaluator, assert_exp_eq):
-    coerce = macro_evaluator._coerce
-    assert coerce(exp.Literal.number(1), int) == 1
-    assert coerce(exp.Literal.number(1.1), float) == 1.1
-    assert coerce(exp.Literal.string("Hi mom"), str) == "Hi mom"
-    assert coerce(exp.true(), bool) is True
-
-    # Coercing a string literal to a column should return a column with the same name
-    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Column), exp.column("order"))
-    # Not possible to coerce this string literal Cast to an exp.Column node -- so it should just return the input
-    assert_exp_eq(
-        coerce(exp.Literal.string("order::date"), exp.Column), exp.Literal.string("order::date")
-    )
-    # This however, is correctly coercible since it's a cast
-    assert_exp_eq(
-        coerce(exp.Literal.string("order::date"), exp.Cast), exp.cast(exp.column("order"), "DATE")
-    )
-
-    # Here we resolve ambiguity via the user type hint
-    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Identifier), exp.to_identifier("order"))
-    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Table), exp.table_("order"))
-
-    # Resolve a union type hint by choosing the first one that works
-    assert_exp_eq(
-        coerce(exp.Literal.string("order::date"), t.Union[exp.Column, exp.Cast]),
-        exp.cast(exp.column("order"), "DATE"),
-    )
-
-    # Simply ask for a string, and always get a string
-    assert coerce(exp.column("order"), str) == "order"
-    assert (
-        coerce(parse_one("SELECT x FROM UNNEST(y) AS x WHERE 1 = 1"), str)
-        == "SELECT x FROM UNNEST(y) AS x WHERE 1 = 1"
-    )
-
-    # From a string literal to a Select should parse the string literal, and the inverse operation works as well
-    assert_exp_eq(
-        coerce(exp.Literal.string("SELECT 1 FROM a"), exp.Select), parse_one("SELECT 1 FROM a")
-    )
-    assert coerce(parse_one("SELECT 1 FROM a"), str) == "SELECT 1 FROM a"
-
-    # Get a list of exp directly instead of an exp.Array
-    assert coerce(parse_one("[1, 2, 3]"), list) == [
-        exp.Literal.number(1),
-        exp.Literal.number(2),
-        exp.Literal.number(3),
-    ]
-
-    # Generics work as well, recursively resolving inner types
-    assert coerce(parse_one("[1, 2, 3]"), t.List[int]) == [1, 2, 3]
-    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, int, float]) == (1, 2, 3.0)
-    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, ...]) == (1, 2, 3)
-    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, str, float]) == (1, "2", 3.0)
-    assert coerce(
-        parse_one("[1, 2, [3]]"), t.Tuple[int, str, t.Union[float, t.Tuple[float, ...]]]
-    ) == (1, "2", (3.0,))
-
-    # Using exp.Expression will always return the input expression
-    assert coerce(parse_one("order", into=exp.Column), exp.Expression) == exp.column("order")
-    assert coerce(exp.Literal.string("OK"), exp.Expression) == exp.Literal.string("OK")
-
-    # Strict flag allows raising errors and is used when recursively coercing expressions
-    # otherwise, in general, we want to be lenient and just warn the user when something is not possible
-    with pytest.raises(SQLMeshError):
-        coerce(exp.Literal.string("order"), exp.Select, strict=True)
-
-    with pytest.raises(SQLMeshError):
-        _ = coerce(
-            exp.Literal.string("order::date"),
-            t.Union[exp.Column, exp.Identifier],
-            strict=True,
-        )
diff --git a/tests/utils/test_metaprogramming.py b/tests/utils/test_metaprogramming.py
index 9b83a24e..a469b576 100644
--- a/tests/utils/test_metaprogramming.py
+++ b/tests/utils/test_metaprogramming.py
@@ -140,19 +140,19 @@ def test_func_globals() -> None:
 def test_normalize_source() -> None:
     assert (
         normalize_source(main_func)
-        == """def main_func(y: int):
+        == """def main_func(y):
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
 
-    def closure(z: int):
+    def closure(z):
         return z + Z
     return closure(y) + other_func(Y)"""
     )
 
     assert (
         normalize_source(other_func)
-        == """def other_func(a: int):
+        == """def other_func(a):
     import sqlglot
     sqlglot.parse_one('1')
     pd.DataFrame([{'x': 1}])
@@ -179,12 +179,12 @@ def test_serialize_env() -> None:
             name="main_func",
             alias="MAIN",
             path="test_metaprogramming.py",
-            payload="""def main_func(y: int):
+            payload="""def main_func(y):
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
 
-    def closure(z: int):
+    def closure(z):
         return z + Z
     return closure(y) + other_func(Y)""",
         ),
@@ -236,7 +236,7 @@ class DataClass:
         "other_func": Executable(
             name="other_func",
             path="test_metaprogramming.py",
-            payload="""def other_func(a: int):
+            payload="""def other_func(a):
     import sqlglot
     sqlglot.parse_one('1')
     pd.DataFrame([{'x': 1}])