Feat: traverse python objects in post-order to support custom decorators (#3762)

Author: georgesittas

sqlmesh/utils/metaprogramming.py

@@ -278,68 +278,74 @@ def build_env(
         name: Name of the object in the env.
         path: The module path to serialize. Other modules will not be walked and treated as imports.
     """
+    # We don't rely on `env` to keep track of visited objects, because it's populated in post-order
+    visited: t.Set[str] = set()
+
+    def walk(obj: t.Any, name: str) -> None:
+        obj_module = inspect.getmodule(obj)
+        if name in visited or (obj_module and obj_module.__name__ == "builtins"):
+            return
+
+        visited.add(name)
+        if name not in env:
+            if hasattr(obj, c.SQLMESH_MACRO):
+                # We only need to add the undecorated code of @macro() functions in env, which
+                # is accessible through the `__wrapped__` attribute added by functools.wraps
+                obj = obj.__wrapped__
+            elif callable(obj) and not isinstance(obj, SERIALIZABLE_CALLABLES):
+                obj = getattr(obj, "__wrapped__", None)
+                name = getattr(obj, "__name__", "")
+
+                # Callable class instances shouldn't be serialized (e.g. tenacity.Retrying).
+                # We still want to walk the callables they decorate, though
+                if not isinstance(obj, SERIALIZABLE_CALLABLES) or name in env:
+                    return
+
+            if (
+                not obj_module
+                or not hasattr(obj_module, "__file__")
+                or not _is_relative_to(obj_module.__file__, path)
+            ):
+                env[name] = obj
+                return
+        elif env[name] != obj:
+            raise SQLMeshError(
+                f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
+            )
 
-    obj_module = inspect.getmodule(obj)
-
-    if obj_module and obj_module.__name__ == "builtins":
-        return
-
-    def walk(obj: t.Any) -> None:
         if inspect.isclass(obj):
             for var in decorator_vars(obj):
                 if obj_module and var in obj_module.__dict__:
-                    build_env(
-                        obj_module.__dict__[var],
-                        env=env,
-                        name=var,
-                        path=path,
-                    )
+                    walk(obj_module.__dict__[var], var)
 
             for base in obj.__bases__:
-                build_env(base, env=env, name=base.__qualname__, path=path)
+                walk(base, base.__qualname__)
 
             for k, v in obj.__dict__.items():
                 if k.startswith("__"):
                     continue
-                # traverse methods in a class to find global references
+
+                # Traverse methods in a class to find global references
                 if isinstance(v, (classmethod, staticmethod)):
                     v = v.__func__
+
                 if callable(v):
-                    # if the method is a part of the object, walk it
-                    # else it is a global function and we just store it
+                    # Walk the method if it's part of the object, else it's a global function and we just store it
                     if v.__qualname__.startswith(obj.__qualname__):
-                        walk(v)
+                        for k, v in func_globals(v).items():
+                            walk(v, k)
                     else:
-                        build_env(v, env=env, name=v.__name__, path=path)
+                        walk(v, v.__name__)
         elif callable(obj):
             for k, v in func_globals(obj).items():
-                build_env(v, env=env, name=k, path=path)
-
-    if name not in env:
-        if hasattr(obj, c.SQLMESH_MACRO):
-            # We only need to add the undecorated code of @macro() functions in env, which
-            # is accessible through the `__wrapped__` attribute added by functools.wraps
-            obj = obj.__wrapped__
-        elif callable(obj) and not isinstance(obj, SERIALIZABLE_CALLABLES):
-            obj = getattr(obj, "__wrapped__", None)
-            name = getattr(obj, "__name__", "")
-
-            # Callable class instances shouldn't be serialized (e.g. tenacity.Retrying).
-            # We still want to walk the callables they decorate, though
-            if not isinstance(obj, SERIALIZABLE_CALLABLES) or name in env:
-                return
+                walk(v, k)
 
+        # We store the object in the environment after its dependencies, because otherwise we
+        # could crash at environment hydration time, since dicts are ordered and the top-level
+        # objects would be loaded before their dependencies.
         env[name] = obj
-        if (
-            obj_module
-            and hasattr(obj_module, "__file__")
-            and _is_relative_to(obj_module.__file__, path)
-        ):
-            walk(env[name])
-    elif env[name] != obj:
-        raise SQLMeshError(
-            f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
-        )
+
+    walk(obj, name)
 
 
 class ExecutableKind(str, Enum):

tests/utils/test_metaprogramming.py

@@ -115,6 +115,18 @@ def fetch_data():
     return "'test data'"
 
 
+def custom_decorator(_func):
+    def wrapper(*args, **kwargs):
+        return _func(*args, **kwargs)
+
+    return wrapper
+
+
+@custom_decorator
+def function_with_custom_decorator():
+    return
+
+
 def main_func(y: int, foo=exp.true(), *, bar=expressions.Literal.number(1) + 2) -> int:
     """DOC STRING"""
     sqlglot.parse_one("1")
@@ -123,6 +135,7 @@ def main_func(y: int, foo=exp.true(), *, bar=expressions.Literal.number(1) + 2)
     noop_metadata()
     normalize_model_name("test")
     fetch_data()
+    function_with_custom_decorator()
 
     def closure(z: int) -> int:
         return z + Z
@@ -147,6 +160,7 @@ def test_func_globals() -> None:
         "expressions": exp,
         "fetch_data": fetch_data,
         "test_context_manager": test_context_manager,
+        "function_with_custom_decorator": function_with_custom_decorator,
     }
     assert func_globals(other_func) == {
         "X": 1,
@@ -180,6 +194,7 @@ def test_normalize_source() -> None:
     noop_metadata()
     normalize_model_name('test')
     fetch_data()
+    function_with_custom_decorator()
 
     def closure(z: int):
         return z + Z
@@ -226,6 +241,7 @@ def test_serialize_env() -> None:
     noop_metadata()
     normalize_model_name('test')
     fetch_data()
+    function_with_custom_decorator()
 
     def closure(z: int):
         return z + Z
@@ -354,4 +370,28 @@ def fetch_data():
             path="test_metaprogramming.py",
             alias="f",
         ),
+        "function_with_custom_decorator": Executable(
+            name="wrapper",
+            path="test_metaprogramming.py",
+            payload="""def wrapper(*args, **kwargs):
+    return _func(*args, **kwargs)""",
+            alias="function_with_custom_decorator",
+        ),
+        "custom_decorator": Executable(
+            name="custom_decorator",
+            path="test_metaprogramming.py",
+            payload="""def custom_decorator(_func):
+
+    def wrapper(*args, **kwargs):
+        return _func(*args, **kwargs)
+    return wrapper""",
+        ),
+        "_func": Executable(
+            name="function_with_custom_decorator",
+            path="test_metaprogramming.py",
+            payload="""@custom_decorator
+def function_with_custom_decorator():
+    return""",
+            alias="_func",
+        ),
     }