Drop cache

libs/langgraph/langgraph/pregel/utils.py

@@ -1,7 +1,10 @@
-from typing import Optional
+import ast
+import inspect
+import textwrap
+from typing import Any, Callable, Optional
 
 from langchain_core.runnables import RunnableLambda, RunnableSequence
-from langchain_core.runnables.utils import get_function_nonlocals
+from typing_extensions import override
 
 from langgraph.checkpoint.base import ChannelVersions
 from langgraph.pregel.protocol import PregelProtocol
@@ -55,3 +58,152 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[PregelProtocol]:
                 )
 
     return None
+
+
+def get_function_nonlocals(func: Callable) -> list[Any]:
+    """Get the nonlocal variables accessed by a function.
+
+    Args:
+        func: The function to check.
+
+    Returns:
+        List[Any]: The nonlocal variables accessed by the function.
+    """
+    try:
+        code = inspect.getsource(func)
+        tree = ast.parse(textwrap.dedent(code))
+        visitor = FunctionNonLocals()
+        visitor.visit(tree)
+        values: list[Any] = []
+        closure = (
+            inspect.getclosurevars(func.__wrapped__)
+            if hasattr(func, "__wrapped__") and callable(func.__wrapped__)
+            else inspect.getclosurevars(func)
+        )
+        candidates = {**closure.globals, **closure.nonlocals}
+        for k, v in candidates.items():
+            if k in visitor.nonlocals:
+                values.append(v)
+            for kk in visitor.nonlocals:
+                if "." in kk and kk.startswith(k):
+                    vv = v
+                    for part in kk.split(".")[1:]:
+                        if vv is None:
+                            break
+                        else:
+                            try:
+                                vv = getattr(vv, part)
+                            except AttributeError:
+                                break
+                    else:
+                        values.append(vv)
+    except (SyntaxError, TypeError, OSError, SystemError):
+        return []
+
+    return values
+
+
+class FunctionNonLocals(ast.NodeVisitor):
+    """Get the nonlocal variables accessed of a function."""
+
+    def __init__(self) -> None:
+        self.nonlocals: set[str] = set()
+
+    @override
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
+        """Visit a function definition.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        visitor = NonLocals()
+        visitor.visit(node)
+        self.nonlocals.update(visitor.loads - visitor.stores)
+
+    @override
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
+        """Visit an async function definition.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        visitor = NonLocals()
+        visitor.visit(node)
+        self.nonlocals.update(visitor.loads - visitor.stores)
+
+    @override
+    def visit_Lambda(self, node: ast.Lambda) -> Any:
+        """Visit a lambda function.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        visitor = NonLocals()
+        visitor.visit(node)
+        self.nonlocals.update(visitor.loads - visitor.stores)
+
+
+class NonLocals(ast.NodeVisitor):
+    """Get nonlocal variables accessed."""
+
+    def __init__(self) -> None:
+        self.loads: set[str] = set()
+        self.stores: set[str] = set()
+
+    @override
+    def visit_Name(self, node: ast.Name) -> Any:
+        """Visit a name node.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        if isinstance(node.ctx, ast.Load):
+            self.loads.add(node.id)
+        elif isinstance(node.ctx, ast.Store):
+            self.stores.add(node.id)
+
+    @override
+    def visit_Attribute(self, node: ast.Attribute) -> Any:
+        """Visit an attribute node.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        if isinstance(node.ctx, ast.Load):
+            parent = node.value
+            attr_expr = node.attr
+            while isinstance(parent, ast.Attribute):
+                attr_expr = parent.attr + "." + attr_expr
+                parent = parent.value
+            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)