Fix!: include names of decorator argument references when building python env

[georgesittas]

Fixes #3640

The main difficulty in solving the linked issue was making SQLMesh detect the stop_after_attempt reference, in order to extract and inject it into the python environment so it can be (de)serialized:

@retry(stop=stop_after_attempt(3))
def fetch_data():
    return "test data"

The decorator names themselves were already being picked up in the previous decorators helper. In this PR I expand that functionality by also searching for the decorator call arguments' references.

I implemented a new visitor class in order to limit the search to only the decorator sub-trees, instead of traversing the whole function tree. Without this, we'd have to do a nested walk within the root node traversal loop, leading to unnecessarily revisiting nodes under the decorator sub-trees.

Other than that, I also made sure to exclude callable instances of classes, because they can't be serialized. One such example is the tenacity.Retrying class (source). Letting these instances into the python environment doesn't "just work" unfortunately, because an error is raised once we reach this section, since there's no __name__ attribute in them.

I verified that this fix works for a project with the model of interest, as well as the example project we have under test_metaprogramming.py.

[izeigerman]

Why is this import included twice?

[georgesittas]

This section brings in the additional "weird" references:

(Pdb) p list(zip(code.co_freevars, func.__closure__))
[('f', <cell at 0x143ae7430: function object at 0x143b1dc60>), ('self', <cell at 0x143ae73a0: Retrying object at 0x143ae6600>), ('wrapped_f', <cell at 0x143ae73d0: function object at 0x143b1e160>)]

This checks out with tenacity's decorator implementation (notice the f and wrapped_f definitions): https://github.com/jd/tenacity/blob/0d40e76f7d06d631fb127e1ec58c8bd776e70d49/tenacity/__init__.py#L322-L346.

So, regarding the two imports:

• The first one corresponds to the ref extracted using _code_globals(code) on the execute function (it brings in the names ['fetch_data', 'pd'])
• The second one corresponds to wrapped_f.

I have an idea on how to fix this.

[georgesittas]

@izeigerman @tobymao thinking out loud: what about only extracting closure variables for functions that are defined within the module path corresponding to the active sqlmesh context? Something along the lines of:

-        if func.__closure__:
-            for var, value in zip(code.co_freevars, func.__closure__):
-                variables[var] = value.cell_contents
+        if func.__closure__ and _is_relative_to(func.__globals__.get("__file__"), path):
+                for var, value in zip(code.co_freevars, func.__closure__):
+                    variables[var] = value.cell_contents
+
+        if hasattr(func, "__wrapped__"):
+            variables.update(func_globals(func.__wrapped__, path=path))

Given these changes I'm seeing two failed tests:

1. y needs to be excluded here
2. 'f', 'func' and 'wrapped_f' need to be excluded from the test_serialize_env test in this PR (functools is still being picked up by decorator_vars above this section)

I added the branch with __wrapped__ because without it the retry and stop_after_attempt dependencies aren't included in the serialized env, which I believe is incorrect?

[izeigerman]

Why do we even capture this import?

[georgesittas]

Because it's included in argument list of fetch_data's decorator:

(Pdb) p decorator_vars(func, root_node=root_node)
['wraps', 'f', 'functools', 'WRAPPER_ASSIGNMENTS']

Due to this line, specifically.

[izeigerman]

Where does his alias come from?

[georgesittas]

I think this reply answers it.