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

Author: georgesittas

.pre-commit-config.yaml

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

setup.cfg

@@ -18,9 +18,6 @@ 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,7 +32,6 @@
     },
     setup_requires=["setuptools_scm"],
     install_requires=[
-        "astor",
         "click",
         "croniter",
         "duckdb!=0.10.3",

sqlmesh/core/model/common.py

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

sqlmesh/utils/metaprogramming.py

@@ -4,7 +4,6 @@
 import dis
 import importlib
 import inspect
-import linecache
 import os
 import re
 import sys
@@ -15,14 +14,17 @@
 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
 
-IGNORE_DECORATORS = {"macro", "model", "signal"}
+IGNORED_DECORATORS = {"macro", "model", "signal"}
+IGNORED_DECORATOR_CALL_REGEX = re.compile(
+    rf'(?s)@({"|".join(IGNORED_DECORATORS)})\s*\(.*?\)(\s|#.*?\n)*def'
+)
+
 SERIALIZABLE_CALLABLES = (type, types.FunctionType)
 LITERALS = (Number, str, bytes, tuple, list, dict, set, bool)
 
@@ -100,40 +102,6 @@ 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] = []
@@ -149,7 +117,7 @@ def _extract_dependencies(self, node: ast.ClassDef | ast.FunctionDef) -> None:
                 else:
                     continue
 
-                if dep in IGNORE_DECORATORS:
+                if dep in IGNORED_DECORATORS:
                     dependencies = []
                     break
 
@@ -166,53 +134,9 @@ 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(getsource(func)))
+    return ast.parse(textwrap.dedent(inspect.getsource(func)))
 
 
 def _decorator_name(decorator: ast.expr) -> str:
@@ -234,33 +158,14 @@ def decorator_vars(func: t.Callable, root_node: t.Optional[ast.Module] = None) -
     return unique(finder.dependencies)
 
 
-def normalize_source(obj: t.Any) -> str:
-    """Rewrites an object's source with formatting and doc strings removed by using Python ast.
-
-    Args:
-        obj: The object to fetch source from and convert to a string.
-
-    Returns:
-        A string representation of the normalized function.
+def remove_ignored_decorators(source: str) -> str:
     """
-    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)
+    Removes decorator calls like @model(...) from the Python source code.
 
-            # 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()
+    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.
+    """
+    return IGNORED_DECORATOR_CALL_REGEX.sub("def", source)
 
 
 def build_env(
@@ -447,7 +352,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=normalize_source(v),
+                    payload=remove_ignored_decorators(inspect.getsource(v).strip()),
                     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

@@ -0,0 +1,11 @@
+# This module is conditionally imported to demonstrate that we can
+# serialize `match` statements that are available only for > v3.9.
+# If we included this code as-is in test_metaprogramming.py, the
+# test harness would crash when using Python 3.9 in CI
+
+
+def match_expression():
+    match 5:
+        case 5:
+            return 1
+    return 0

tests/core/test_model.py

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