Support lambda extraction by AST parsing

Author: Евгений Блинов

getsources/clear.py

@@ -1,11 +1,33 @@
+from ast import Lambda, get_source_segment, parse, walk
 from typing import Any, Callable
 
 from getsources import getsource
+from getsources.errors import UncertaintyWithLambdasError
+from getsources.helpers.is_lambda import is_lambda
 
 
 def getclearsource(function: Callable[..., Any]) -> str:
     source_code = getsource(function)
 
+    if is_lambda(function):
+        stripped_source_code = source_code.strip()
+        tree = parse(stripped_source_code)
+
+        first = True
+        lambda_node = None
+        for node in walk(tree):
+            if isinstance(node, Lambda):
+                if not first:
+                    raise UncertaintyWithLambdasError('Several lambda functions are defined in a single line of code, can\'t pick the one.')
+                lambda_node = node
+                first = False
+
+        segment_source = get_source_segment(stripped_source_code, lambda_node)
+        if segment_source is None:
+            raise UncertaintyWithLambdasError('It seems that the AST for the lambda function has been modified; can\'t extract the source code.')
+        return segment_source
+
+
     splitted_source_code = source_code.split('\n')
 
     indent = 0

getsources/errors.py

@@ -0,0 +1,2 @@
+class UncertaintyWithLambdasError(Exception):
+    ...

getsources/hash.py

@@ -1,6 +1,6 @@
 import hashlib
+from ast import Constant, Expr, Lambda, get_source_segment, parse, walk
 from typing import Any, Callable
-from ast import parse, Expr, Constant, Lambda, get_source_segment, walk
 
 from getsources import getclearsource
 from getsources.helpers.is_lambda import is_lambda

tests/test_clear.py

@@ -49,8 +49,8 @@ def function_2(a, b):
 def test_lambda():
     function = lambda x: x
 
-    assert getclearsource(function) == 'function = lambda x: x'
-    assert getclearsource(global_function_3) == 'global_function_3 = lambda x: x'
+    assert getclearsource(function) == 'lambda x: x'
+    assert getclearsource(global_function_3) == 'lambda x: x'
 
 
 def test_usual_methods():
@@ -171,4 +171,4 @@ def test_lambda_in_REPL():  # noqa: N802
 
     child.sendline("exit()")
 
-    assert any('function = lambda x: x' in x for x in after)
+    assert any('lambda x: x' in x for x in after)