core: Add support for calls in get_function_nonlocals

langchain-ai · Jan 16, 2025 · 13347e1 · 13347e1
1 parent 667d2a5
commit 13347e1
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/libs/core/langchain_core/runnables/utils.py b/libs/core/langchain_core/runnables/utils.py
@@ -271,6 +271,20 @@ def visit_Attribute(self, node: ast.Attribute) -> Any:
             if isinstance(parent, ast.Name):
                 self.loads.add(parent.id + "." + attr_expr)
                 self.loads.discard(parent.id)
+            elif isinstance(parent, ast.Call):
+                if isinstance(parent.func, ast.Name):
+                    self.loads.add(parent.func.id)
+                else:
+                    parent = parent.func
+                    attr_expr = ""
+                    while isinstance(parent, ast.Attribute):
+                        if attr_expr:
+                            attr_expr = parent.attr + "." + attr_expr
+                        else:
+                            attr_expr = parent.attr
+                        parent = parent.value
+                    if isinstance(parent, ast.Name):
+                        self.loads.add(parent.id + "." + attr_expr)
 
 
 class FunctionNonLocals(ast.NodeVisitor):