Fix: avoid redundant object traversals when building python envs (#4295)

Author: xardasos

sqlmesh/utils/metaprogramming.py

@@ -338,43 +338,43 @@ def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> None:
             ):
                 env[name] = (obj, is_metadata)
                 return
+
+            if inspect.isclass(obj):
+                for var in decorator_vars(obj):
+                    if obj_module and var in obj_module.__dict__:
+                        walk(obj_module.__dict__[var], var, is_metadata)
+
+                for base in obj.__bases__:
+                    walk(base, base.__qualname__, is_metadata)
+
+                for k, v in obj.__dict__.items():
+                    if k.startswith("__"):
+                        continue
+
+                    # Traverse methods in a class to find global references
+                    if isinstance(v, (classmethod, staticmethod)):
+                        v = v.__func__
+
+                    if callable(v):
+                        # 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__):
+                            for k, v in func_globals(v).items():
+                                walk(v, k, is_metadata)
+                        else:
+                            walk(v, v.__name__, is_metadata)
+            elif callable(obj):
+                for k, v in func_globals(obj).items():
+                    walk(v, k, is_metadata)
+
+            # 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, is_metadata)
         elif not _globals_match(env[name][0], obj):
             raise SQLMeshError(
                 f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
             )
 
-        if inspect.isclass(obj):
-            for var in decorator_vars(obj):
-                if obj_module and var in obj_module.__dict__:
-                    walk(obj_module.__dict__[var], var, is_metadata)
-
-            for base in obj.__bases__:
-                walk(base, base.__qualname__, is_metadata)
-
-            for k, v in obj.__dict__.items():
-                if k.startswith("__"):
-                    continue
-
-                # Traverse methods in a class to find global references
-                if isinstance(v, (classmethod, staticmethod)):
-                    v = v.__func__
-
-                if callable(v):
-                    # 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__):
-                        for k, v in func_globals(v).items():
-                            walk(v, k, is_metadata)
-                    else:
-                        walk(v, v.__name__, is_metadata)
-        elif callable(obj):
-            for k, v in func_globals(obj).items():
-                walk(v, k, is_metadata)
-
-        # 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, is_metadata)
-
     # The "metadata only" annotation of the object is transitive
     walk(obj, name, is_metadata_obj or getattr(obj, c.SQLMESH_METADATA, None))
 

tests/utils/test_metaprogramming.py

@@ -17,6 +17,7 @@
 import tests.utils.test_date as test_date
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core import constants as c
+from sqlmesh.core.macros import RuntimeStage
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.metaprogramming import (
     Executable,
@@ -47,7 +48,7 @@ def test_print_exception(mocker: MockerFixture):
     except Exception as ex:
         print_exception(ex, test_env, out_mock)
 
-    expected_message = r"""  File ".*?/tests/utils/test_metaprogramming\.py", line 46, in test_print_exception
+    expected_message = r"""  File ".*?/tests/utils/test_metaprogramming\.py", line 47, in test_print_exception
     eval\("test_fun\(\)", env\)
 
   File "<string>", line 1, in <module>
@@ -140,6 +141,18 @@ def closure(z: int) -> int:
     return closure(y) + other_func(Y)
 
 
+def macro1() -> str:
+    print("macro1 hello there")
+    print(RuntimeStage.CREATING)
+    return "1"
+
+
+def macro2() -> str:
+    print("macro2 hello there")
+    print(RuntimeStage.LOADING)
+    return "2"
+
+
 def test_func_globals() -> None:
     assert func_globals(main_func) == {
         "Y": 2,
@@ -405,3 +418,38 @@ def function_with_custom_decorator():
     # Every object is treated as "metadata only", transitively
     assert all(is_metadata for (_, is_metadata) in env.values())
     assert serialized_env == expected_env
+
+
+def test_serialize_env_with_enum_import_appearing_in_two_functions() -> None:
+    path = Path("tests/utils")
+    env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
+
+    build_env(macro1, env=env, name="macro1", path=path)
+    build_env(macro2, env=env, name="macro2", path=path)
+
+    serialized_env = serialize_env(env, path=path)  # type: ignore
+    assert prepare_env(serialized_env)
+
+    expected_env = {
+        "RuntimeStage": Executable(
+            payload="from sqlmesh.core.macros import RuntimeStage", kind=ExecutableKind.IMPORT
+        ),
+        "macro1": Executable(
+            payload="""def macro1():
+    print('macro1 hello there')
+    print(RuntimeStage.CREATING)
+    return '1'""",
+            name="macro1",
+            path="test_metaprogramming.py",
+        ),
+        "macro2": Executable(
+            payload="""def macro2():
+    print('macro2 hello there')
+    print(RuntimeStage.LOADING)
+    return '2'""",
+            name="macro2",
+            path="test_metaprogramming.py",
+        ),
+    }
+
+    assert serialized_env == expected_env