More robust parsing of ast's from callables (#97)

* Better testing for edge cases
* Deal with indents more naturally
* Allow supplying method name to parser for cases where things are too deep

Author: gordonwatts

func_adl/util_ast.py

@@ -3,7 +3,7 @@
 import ast
 import inspect
 import sys
-from typing import Any, Callable, Dict, List, Optional, Union, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 
 # Some functions to enable backwards compatibility.
 # Capability may be degraded in older versions - particularly 3.6.
@@ -351,7 +351,25 @@ def global_getclosurevars(f: Callable) -> inspect.ClosureVars:
     return cv
 
 
-def parse_as_ast(ast_source: Union[str, ast.AST, Callable]) -> ast.Lambda:
+def _realign_indent(s: str) -> str:
+    """Move the first line to be at zero indent, and then apply same for everything
+    below.
+
+    Args:
+        s (str): The string with indents
+
+    Returns:
+        str: Unindent string
+    """
+    lines = s.split("\n")
+    spaces = len(lines[0]) - len(lines[0].lstrip())
+    stripped_lines = [ln[spaces:] for ln in lines]
+    return "\n".join(stripped_lines)
+
+
+def parse_as_ast(
+    ast_source: Union[str, ast.AST, Callable], caller_name: Optional[str] = None
+) -> ast.Lambda:
     r"""Return an AST for a lambda function from several sources.
 
     We are handed one of several things:
@@ -363,19 +381,25 @@ def parse_as_ast(ast_source: Union[str, ast.AST, Callable]) -> ast.Lambda:
 
     Args:
         ast_source:     An AST or text string that represnets the lambda.
+        caller_name:    The name of the function that the lambda is an arg to. If it
+            is none, then it will attempt to scan the stack frame above to figure it out.
 
     Returns:
         An ast starting from the Lambda AST node.
     """
     if callable(ast_source):
-        source = inspect.getsource(ast_source).strip()
-
-        # Look for the name of the calling function (e.g. 'Select' or 'Where')
-        caller_name = inspect.currentframe().f_back.f_code.co_name  # type: ignore
-        caller_idx = source.find(caller_name)
-        # If found, parse the string between the parentheses of the function call
-        if caller_idx > -1:
-            source = source[caller_idx + len(caller_name) :]
+        source = _realign_indent(inspect.getsource(ast_source))
+
+        def find_next_lambda(method_name: str, source: str) -> Tuple[Optional[str], str]:
+            "Find the lambda starting at the name"
+            caller_idx = source.find(method_name)
+
+            # If we couldn't find it, then we need to parse the whole thing.
+            if caller_idx == -1:
+                return None, source
+
+            source = source[caller_idx + len(method_name) :]
+
             i = 0
             open_count = 0
             while True:
@@ -387,42 +411,81 @@ def parse_as_ast(ast_source: Union[str, ast.AST, Callable]) -> ast.Lambda:
                 if open_count == 0:
                     break
                 i += 1
-            stem = source[i + 1 :]
-            new_line = stem.find("\n")
-            next_caller = stem.find(caller_name)
-            if next_caller > -1 and (new_line < 0 or new_line > next_caller):
-                raise ValueError(
-                    f"Found two calls to {caller_name} on same line - " "split accross lines"
-                )
-            source = source[: i + 1]
-
-        def parse(src: str) -> Optional[ast.Module]:
-            try:
-                return ast.parse(src)
-            except SyntaxError:
-                return None
-
-        # Special case ending with a close parenthesis at the end of a line.
-        src_ast = parse(source)
-        if not src_ast and source.endswith(")"):
-            src_ast = parse(source[:-1])
 
-        if not src_ast:
-            raise ValueError(f"Unable to recover source for function {ast_source}.")
+            lambda_source = source[: i + 1]
+            remaining_source = source[i + 1 :]
+
+            return lambda_source, remaining_source
+
+        # Look for the name of the calling function (e.g. 'Select' or 'Where', etc.) and
+        # find all the instances on this line.
+        if caller_name is None:
+            caller_name = inspect.currentframe().f_back.f_code.co_name  # type: ignore
+
+        found_lambdas: List[str] = []
+        while True:
+            lambda_source, remaining_source = find_next_lambda(caller_name, source)
+            if lambda_source is None:
+                break
+            source = remaining_source
+            found_lambdas.append(lambda_source)
 
-        # If this is a function, not a lambda, then we can morph and return that.
-        if len(src_ast.body) == 1 and isinstance(src_ast.body[0], ast.FunctionDef):
-            lda = rewrite_func_as_lambda(src_ast.body[0])  # type: ignore
+        if len(found_lambdas) == 0:
+            found_lambdas.append(source)
+
+        # Parse them as a lambda function
+        def parse(src: str) -> Optional[ast.Lambda]:
+            while True:
+                try:
+                    a_module = ast.parse(src)
+                    # If this is a function, not a lambda, then we can morph and return that.
+                    if len(a_module.body) == 1 and isinstance(a_module.body[0], ast.FunctionDef):
+                        lda = rewrite_func_as_lambda(a_module.body[0])  # type: ignore
+                    else:
+                        lda = next(
+                            (node for node in ast.walk(a_module) if isinstance(node, ast.Lambda)),
+                            None,
+                        )
+
+                    if lda is None:
+                        raise ValueError(
+                            f"Unable to recover source for function {ast_source} - '{src}'."
+                        )
+                    return lda
+                except SyntaxError:
+                    pass
+                if src.endswith(")"):
+                    src = src[:-1]
+                else:
+                    return None
+
+        parsed_lambdas = [parse(src) for src in found_lambdas]
+
+        # If we have more than one lambda, there are some tricks we can try - like argument names,
+        # to see if they are different.
+        src_ast: Optional[ast.Lambda] = None
+        if len(found_lambdas) > 1:
+            caller_arg_list = inspect.getfullargspec(ast_source).args
+            for idx, p_lambda in enumerate(parsed_lambdas):
+                lambda_args = [a.arg for a in p_lambda.args.args]  # type: ignore
+                if lambda_args == caller_arg_list:
+                    if src_ast is not None:
+                        raise ValueError(
+                            f"Found two calls to {caller_name} on same line - "
+                            "split accross lines or change lambda argument names so they "
+                            "are different."
+                        )
+                    src_ast = p_lambda
         else:
-            lda = next((node for node in ast.walk(src_ast) if isinstance(node, ast.Lambda)), None)
+            assert len(found_lambdas) == 1
+            src_ast = parsed_lambdas[0]
 
-            if lda is None:
-                raise ValueError(f"Unable to recover source for function {ast_source}.")
+        if not src_ast:
+            raise ValueError(f"Unable to recover source for function {ast_source}.")
 
         # Since this is a function in python, we can look for lambda capture.
         call_args = global_getclosurevars(ast_source)
-
-        return _rewrite_captured_vars(call_args).visit(lda)
+        return _rewrite_captured_vars(call_args).visit(src_ast)
 
     elif isinstance(ast_source, str):
         a = ast.parse(ast_source.strip())  # type: ignore

tests/test_util_ast.py

@@ -5,6 +5,7 @@
 import pytest
 
 from func_adl.util_ast import (
+    _realign_indent,
     as_ast,
     function_call,
     lambda_args,
@@ -366,28 +367,92 @@ def test_known_global_function():
     assert "Name(id='global_doit_non_func'" in ast.dump(f)
 
 
-# TODO: This test is not compatible with the black formatter - it puts
-#       the doit call on the line twice.
-# def test_parse_continues():
-#     "Emulate the syntax you often find when you have a multistep query"
-#     found = []
+def test_parse_continues():
+    "Emulate the syntax you often find when you have a multistep query"
+    found = []
+
+    class my_obj:
+        def do_it(self, x: Callable):
+            found.append(parse_as_ast(x))
+            return self
+
+    (my_obj().do_it(lambda x: x + 1).do_it(lambda y: y * 2))
+
+    assert len(found) == 2
+    l1, l2 = found
+    assert isinstance(l1, ast.Lambda)
+    assert isinstance(l1.body, ast.BinOp)
+    assert isinstance(l1.body.op, ast.Add)
+
+    assert isinstance(l2, ast.Lambda)
+    assert isinstance(l2.body, ast.BinOp)
+    assert isinstance(l2.body.op, ast.Mult)
+
+
+def test_decorator_parse():
+    "More general case"
+
+    seen_lambdas = []
+
+    def dec_func(x: Callable):
+        def make_it(y: Callable):
+            return y
+
+        seen_lambdas.append(parse_as_ast(x))
+        return make_it
+
+    @dec_func(lambda y: y + 2)
+    def doit(x):
+        return x + 1
+
+    assert len(seen_lambdas) == 1
+    l1 = seen_lambdas[0]
+    assert isinstance(l1.body, ast.BinOp)
+    assert isinstance(l1.body.op, ast.Add)
+
 
-#     class my_obj:
-#         def do_it(self, x: Callable):
-#             found.append(parse_as_ast(x))
-#             return self
+def test_indent_parse():
+    "More general case"
 
-#     (my_obj().do_it(lambda x: x + 1).do_it(lambda y: y * 2))
+    seen_funcs = []
 
-#     assert len(found) == 2
-#     l1, l2 = found
-#     assert isinstance(l1, ast.Lambda)
-#     assert isinstance(l1.body, ast.BinOp)
-#     assert isinstance(l1.body.op, ast.Add)
+    class h:
+        @staticmethod
+        def dec_func(x: Callable):
+            def make_it(y: Callable):
+                return y
 
-#     assert isinstance(l2, ast.Lambda)
-#     assert isinstance(l2.body, ast.BinOp)
-#     assert isinstance(l2.body.op, ast.Mult)
+            seen_funcs.append(x)
+            return make_it
+
+    class yo_baby:
+        @h.dec_func(lambda y: y + 2)
+        def doit(self, x: int):
+            ...
+
+    assert len(seen_funcs) == 1
+    l1 = parse_as_ast(seen_funcs[0], "dec_func")
+    assert isinstance(l1.body, ast.BinOp)
+    assert isinstance(l1.body.op, ast.Add)
+
+
+def test_two_deep_parse():
+    "More general case"
+
+    seen_lambdas = []
+
+    def func_bottom(x: Callable):
+        seen_lambdas.append(parse_as_ast(x))
+
+    def func_top(x: Callable):
+        func_bottom(x)
+
+    func_top(lambda x: x + 1)
+
+    assert len(seen_lambdas) == 1
+    l1 = seen_lambdas[0]
+    assert isinstance(l1.body, ast.BinOp)
+    assert isinstance(l1.body.op, ast.Add)
 
 
 def test_parse_continues_one_line():
@@ -400,7 +465,7 @@ def do_it(self, x: Callable):
             return self
 
     with pytest.raises(Exception) as e:
-        my_obj().do_it(lambda x: x + 1).do_it(lambda y: y * 2)
+        my_obj().do_it(lambda x: x + 1).do_it(lambda x: x * 2)
 
     assert "two" in str(e.value)
 
@@ -416,3 +481,23 @@ def callback(a: ast.arg):
 
     assert recoreded is not None
     assert 22 == ast.literal_eval(recoreded)
+
+
+def test_realign_no_indent():
+    assert _realign_indent("test") == "test"
+
+
+def test_realign_indent_sp():
+    assert _realign_indent("    test") == "test"
+
+
+def test_realign_indent_tab():
+    assert _realign_indent("\ttest") == "test"
+
+
+def test_realign_indent_2lines():
+    assert _realign_indent("    test()\n    dude()") == "test()\ndude()"
+
+
+def test_realign_indent_2lines_uneven():
+    assert _realign_indent("    test()\n        dude()") == "test()\n    dude()"