Revert "Fix!: use inspect.getsource in favor of custom parsing, stop relying on astor" (#3909)

Author: georgesittas

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
         "types": [python]
         files: *files
         require_serial: true
-        exclude: ^(tests/fixtures|tests/core/metaprogramming_test_helper\.py)
+        exclude: ^(tests/fixtures)
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: "fc26039"
     hooks:

setup.cfg

@@ -18,6 +18,9 @@ ignore_missing_imports = True
 [mypy-tests.*]
 disallow_untyped_defs = False
 
+[mypy-astor.*]
+ignore_missing_imports = True
+
 [mypy-IPython.*]
 ignore_missing_imports = True
 

setup.py

@@ -32,6 +32,7 @@
     },
     setup_requires=["setuptools_scm"],
     install_requires=[
+        "astor",
         "click",
         "croniter",
         "duckdb!=0.10.3",

sqlmesh/core/model/common.py

@@ -4,6 +4,7 @@
 import typing as t
 from pathlib import Path
 
+from astor import to_source
 from sqlglot import exp
 from sqlglot.helper import ensure_list
 
@@ -140,7 +141,7 @@ def parse_dependencies(
 
                 def get_first_arg(keyword_arg_name: str) -> t.Any:
                     if node.args:
-                        first_arg: t.Optional[ast.AST] = node.args[0]
+                        first_arg: t.Optional[ast.expr] = node.args[0]
                     else:
                         first_arg = next(
                             (
@@ -151,20 +152,14 @@ def get_first_arg(keyword_arg_name: str) -> t.Any:
                             None,
                         )
 
-                    if first_arg is None:
-                        raise_config_error(
-                            f"Missing {keyword_arg_name} argument in '{func.attr}' call.",
-                            executable.path,
-                        )
-
                     try:
-                        expression = ast.unparse(t.cast(ast.AST, first_arg))
+                        expression = to_source(first_arg)
                         return eval(expression, env)
                     except Exception:
                         if strict_resolution:
-                            raise_config_error(
-                                f"Argument '{expression.strip()}' must be resolvable at parse time",
-                                executable.path,
+                            raise ConfigError(
+                                f"Error resolving dependencies for '{executable.path}'. "
+                                f"Argument '{expression.strip()}' must be resolvable at parse time."
                             )
 
                 if func.value.id == "context" and func.attr in ("table", "resolve_table"):

sqlmesh/utils/metaprogramming.py

@@ -4,6 +4,7 @@
 import dis
 import importlib
 import inspect
+import linecache
 import os
 import re
 import sys
@@ -14,17 +15,14 @@
 from numbers import Number
 from pathlib import Path
 
+from astor import to_source
 
 from sqlmesh.core import constants as c
 from sqlmesh.utils import format_exception, unique
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import PydanticModel
 
-IGNORED_DECORATORS = {"macro", "model", "signal"}
-IGNORED_DECORATOR_CALL_REGEX = re.compile(
-    rf'(?s)@({"|".join(IGNORED_DECORATORS)})\s*\(.*?\)(\s|#.*?\n)*def'
-)
-
+IGNORE_DECORATORS = {"macro", "model", "signal"}
 SERIALIZABLE_CALLABLES = (type, types.FunctionType)
 LITERALS = (Number, str, bytes, tuple, list, dict, set, bool)
 
@@ -102,6 +100,40 @@ def func_globals(func: t.Callable) -> t.Dict[str, t.Any]:
     return variables
 
 
+class ClassFoundException(Exception):
+    pass
+
+
+class _ClassFinder(ast.NodeVisitor):
+    def __init__(self, qualname: str) -> None:
+        self.stack: t.List[str] = []
+        self.qualname = qualname
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        self.stack.append(node.name)
+        self.stack.append("<locals>")
+        self.generic_visit(node)
+        self.stack.pop()
+        self.stack.pop()
+
+    visit_AsyncFunctionDef = visit_FunctionDef  # type: ignore
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        self.stack.append(node.name)
+        if self.qualname == ".".join(self.stack):
+            # Return the decorator for the class if present
+            if node.decorator_list:
+                line_number = node.decorator_list[0].lineno
+            else:
+                line_number = node.lineno
+
+            # decrement by one since lines starts with indexing by zero
+            line_number -= 1
+            raise ClassFoundException(line_number)
+        self.generic_visit(node)
+        self.stack.pop()
+
+
 class _DecoratorDependencyFinder(ast.NodeVisitor):
     def __init__(self) -> None:
         self.dependencies: t.List[str] = []
@@ -117,7 +149,7 @@ def _extract_dependencies(self, node: ast.ClassDef | ast.FunctionDef) -> None:
                 else:
                     continue
 
-                if dep in IGNORED_DECORATORS:
+                if dep in IGNORE_DECORATORS:
                     dependencies = []
                     break
 
@@ -134,9 +166,53 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None:
     visit_AsyncFunctionDef = visit_FunctionDef  # type: ignore
 
 
+def getsource(obj: t.Any) -> str:
+    """Get the source of a function or class.
+
+    inspect.getsource doesn't find decorators in python < 3.9
+    https://github.com/python/cpython/commit/696136b993e11b37c4f34d729a0375e5ad544ade
+    """
+    path = inspect.getsourcefile(obj)
+    if path:
+        module = inspect.getmodule(obj, path)
+
+        if module:
+            lines = linecache.getlines(path, module.__dict__)
+        else:
+            lines = linecache.getlines(path)
+
+        def join_source(lnum: int) -> str:
+            return "".join(inspect.getblock(lines[lnum:]))
+
+        if inspect.isclass(obj):
+            qualname = obj.__qualname__
+            source = "".join(lines)
+            tree = ast.parse(source)
+            class_finder = _ClassFinder(qualname)
+            try:
+                class_finder.visit(tree)
+            except ClassFoundException as e:
+                return join_source(e.args[0])
+        elif inspect.isfunction(obj):
+            obj = obj.__code__
+            if hasattr(obj, "co_firstlineno"):
+                lnum = obj.co_firstlineno - 1
+                pat = re.compile(r"^(\s*def\s)|(\s*async\s+def\s)|(.*(?<!\w)lambda(:|\s))|^(\s*@)")
+                while lnum > 0:
+                    try:
+                        line = lines[lnum]
+                    except IndexError:
+                        raise OSError("lineno is out of bounds")
+                    if pat.match(line):
+                        break
+                    lnum = lnum - 1
+                return join_source(lnum)
+    raise SQLMeshError(f"Cannot find source for {obj}")
+
+
 def parse_source(func: t.Callable) -> ast.Module:
     """Parse a function and returns an ast node."""
-    return ast.parse(textwrap.dedent(inspect.getsource(func)))
+    return ast.parse(textwrap.dedent(getsource(func)))
 
 
 def _decorator_name(decorator: ast.expr) -> str:
@@ -158,14 +234,33 @@ def decorator_vars(func: t.Callable, root_node: t.Optional[ast.Module] = None) -
     return unique(finder.dependencies)
 
 
-def remove_ignored_decorators(source: str) -> str:
-    """
-    Removes decorator calls like @model(...) from the Python source code.
+def normalize_source(obj: t.Any) -> str:
+    """Rewrites an object's source with formatting and doc strings removed by using Python ast.
 
-    We do this because we don't need to serialize the decorator or any value within its argument
-    list when hydrating the python environment; we only need the function definition itself.
+    Args:
+        obj: The object to fetch source from and convert to a string.
+
+    Returns:
+        A string representation of the normalized function.
     """
-    return IGNORED_DECORATOR_CALL_REGEX.sub("def", source)
+    root_node = parse_source(obj)
+
+    for node in ast.walk(root_node):
+        if isinstance(node, (ast.FunctionDef, ast.ClassDef)):
+            for decorator in node.decorator_list:
+                if _decorator_name(decorator) in IGNORE_DECORATORS:
+                    node.decorator_list.remove(decorator)
+
+            # remove docstrings
+            body = node.body
+            if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Str):
+                node.body = body[1:]
+
+            # remove function return type annotation
+            if isinstance(node, ast.FunctionDef):
+                node.returns = None
+
+    return to_source(root_node).strip()
 
 
 def build_env(
@@ -352,7 +447,7 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
             if relative_obj_file_path:
                 serialized[k] = Executable(
                     name=name,
-                    payload=remove_ignored_decorators(inspect.getsource(v).strip()),
+                    payload=normalize_source(v),
                     kind=ExecutableKind.DEFINITION,
                     # Do `as_posix` to serialize windows path back to POSIX
                     path=t.cast(Path, file_path).relative_to(path.absolute()).as_posix(),

tests/core/metaprogramming_test_helper.py

@@ -4576,7 +4576,7 @@ def test_default_catalog_sql(assert_exp_eq):
 
 
 def test_default_catalog_python():
-    HASH_WITH_CATALOG = "1357281693"
+    HASH_WITH_CATALOG = "770057346"
 
     @model(name="db.table", kind="full", columns={'"COL"': "int"})
     def my_model(context, **kwargs):