This commit is contained in:
Nuno Campos 2023-09-28 17:12:07 +01:00
parent 151f27d502
commit 5ca461160b
2 changed files with 9 additions and 3 deletions

View File

@ -90,8 +90,10 @@ class IsFunctionArgDict(ast.NodeVisitor):
class GetLambdaSource(ast.NodeVisitor): class GetLambdaSource(ast.NodeVisitor):
def __init__(self) -> None: def __init__(self) -> None:
self.source: Optional[str] = None self.source: Optional[str] = None
self.count = 0
def visit_Lambda(self, node: ast.Lambda) -> Any: def visit_Lambda(self, node: ast.Lambda) -> Any:
self.count += 1
if hasattr(ast, "unparse"): if hasattr(ast, "unparse"):
self.source = ast.unparse(node) self.source = ast.unparse(node)
@ -118,10 +120,11 @@ def get_lambda_source(func: Callable) -> Optional[str]:
""" """
try: try:
code = inspect.getsource(func) code = inspect.getsource(func)
print(code)
tree = ast.parse(textwrap.dedent(code)) tree = ast.parse(textwrap.dedent(code))
visitor = GetLambdaSource() visitor = GetLambdaSource()
visitor.visit(tree) visitor.visit(tree)
return visitor.source return visitor.source if visitor.count == 1 else None
except (SyntaxError, TypeError, OSError): except (SyntaxError, TypeError, OSError):
return None return None

View File

@ -2741,6 +2741,9 @@ async def test_runnable_branch_abatch() -> None:
assert await branch.abatch([1, 10, 0]) == [2, 100, -1] assert await branch.abatch([1, 10, 0]) == [2, 100, -1]
@pytest.mark.skipif(
sys.version_info < (3, 9), reason="Requires python version >= 3.9 to run."
)
def test_reprsentation_of_runnables() -> None: def test_reprsentation_of_runnables() -> None:
"""Test representation of runnables.""" """Test representation of runnables."""
runnable = RunnableLambda(lambda x: x * 2) runnable = RunnableLambda(lambda x: x * 2)
@ -2765,9 +2768,9 @@ def test_reprsentation_of_runnables() -> None:
"b": RunnableLambda(lambda x: x * 3), "b": RunnableLambda(lambda x: x * 3),
} }
) == ( ) == (
"RunnableLambda(lambda x: x * 3)\n" "RunnableLambda(...)\n"
"| {\n" "| {\n"
" a: RunnableLambda(...),\n" " a: RunnableLambda(...),\n"
" b: RunnableLambda(...)\n" " b: RunnableLambda(...)\n"
" }" " }"
) ), "repr where code string contains multiple lambdas gives up"