Detect local arguments and functions used in a function

flowrep/models/parsers/dependency_parser.py

@@ -1,84 +1,18 @@
+from __future__ import annotations
+
 import ast
-import types
+import builtins
+import inspect
+import textwrap
 from collections.abc import Callable
 
 from pyiron_snippets import versions
 
-from flowrep.models.parsers import object_scope, parser_helpers
+from flowrep.models.parsers import object_scope
 
 CallDependencies = dict[versions.VersionInfo, Callable]
 
 
-def get_call_dependencies(
-    func: types.FunctionType,
-    version_scraping: versions.VersionScrapingMap | None = None,
-    _call_dependencies: CallDependencies | None = None,
-    _visited: set[str] | None = None,
-) -> CallDependencies:
-    """
-    Recursively collect all callable dependencies of *func* via AST introspection.
-
-    Each dependency is keyed by its :class:`~pyiron_snippets.versions.VersionInfo`
-    and maps to the callables instance with that identity.  The search is depth-first:
-    for every resolved callee that is a :class:`~types.FunctionType` (i.e. has
-    inspectable source), the function recurses into the callee's own scope.
-
-    Args:
-        func: The function whose call-graph to analyse.
-        version_scraping (VersionScrapingMap | None): Since some modules may store
-            their version in other ways, this provides an optional map between module
-            names and callables to leverage for extracting that module's version.
-        _call_dependencies: Accumulator for recursive calls — do not pass manually.
-        _visited: Fully-qualified names already traversed — do not pass manually.
-
-    Returns:
-        A mapping from :class:`VersionInfo` to the callables found under that
-        identity across the entire (sub-)tree.
-    """
-    call_dependencies: CallDependencies = _call_dependencies or {}
-    visited: set[str] = _visited or set()
-
-    func_fqn = versions.VersionInfo.of(func).fully_qualified_name
-    if func_fqn in visited:
-        return call_dependencies
-    visited.add(func_fqn)
-
-    scope = object_scope.get_scope(func)
-    tree = parser_helpers.get_ast_function_node(func)
-    collector = CallCollector()
-    collector.visit(tree)
-
-    for call in collector.calls:
-        try:
-            caller = object_scope.resolve_symbol_to_object(call, scope)
-        except (ValueError, TypeError):
-            continue
-
-        if not callable(caller):  # pragma: no cover
-            # Under remotely normal circumstances, this should be unreachable
-            raise TypeError(
-                f"Caller {caller} is not callable, yet was generated from the list of "
-                f"ast.Call calls, in particular {call}. We're expecting these to "
-                f"actually connect to callables. Please raise a GitHub issue if you "
-                f"think this is not a mistake."
-            )
-
-        info = versions.VersionInfo.of(caller, version_scraping=version_scraping)
-        # In principle, we open ourselves to overwriting an existing dependency here,
-        # but it would need to somehow have exactly the same version info (including
-        # qualname) yet be a different object.
-        # This ought not happen by accident, and in case it somehow does happen on
-        # purpose (it probably shouldn't), we just silently keep the more recent one.
-
-        call_dependencies[info] = caller
-
-        # Depth-first search on dependencies — only possible when we have source
-        if isinstance(caller, types.FunctionType):
-            get_call_dependencies(caller, version_scraping, call_dependencies, visited)
-
-    return call_dependencies
-
-
 def split_by_version_availability(
     call_dependencies: CallDependencies,
 ) -> tuple[CallDependencies, CallDependencies]:
@@ -102,10 +36,61 @@ def split_by_version_availability(
     return has_version, no_version
 
 
-class CallCollector(ast.NodeVisitor):
+class UndefinedVariableVisitor(ast.NodeVisitor):
     def __init__(self):
-        self.calls: list[ast.expr] = []
-
-    def visit_Call(self, node: ast.Call) -> None:
-        self.calls.append(node.func)
+        self.used_vars = set()
+        self.defined_vars = set()
+
+    def visit_Name(self, node):
+        if isinstance(node.ctx, ast.Load):  # Variable is being used
+            self.used_vars.add(node.id)
+        elif isinstance(node.ctx, ast.Store):  # Variable is being defined
+            self.defined_vars.add(node.id)
+
+    def visit_FunctionDef(self, node):
+        # Add function arguments to defined variables
+        for arg in node.args.args:
+            self.defined_vars.add(arg.arg)
         self.generic_visit(node)
+
+
+def find_undefined_variables(
+    func_or_var: Callable | object,
+) -> set[str]:
+    if callable(func_or_var):
+        source = textwrap.dedent(inspect.getsource(func_or_var))
+    else:
+        source = str(func_or_var)
+    tree = ast.parse(source)
+
+    visitor = UndefinedVariableVisitor()
+    visitor.visit(tree)
+    undefined_vars = visitor.used_vars - visitor.defined_vars
+    return undefined_vars.difference(set(dir(builtins)))
+
+
+def get_call_dependencies(
+    func_or_var: Callable | object,
+    version_scraping: versions.VersionScrapingMap | None = None,
+    _call_dependencies: CallDependencies | None = None,
+    _visited: set[str] | None = None,
+) -> CallDependencies:
+
+    call_dependencies: CallDependencies = _call_dependencies or {}
+    visited: set[str] = _visited or set()
+
+    func_fqn = versions.VersionInfo.of(func_or_var).fully_qualified_name
+    if func_fqn in visited:
+        return call_dependencies
+    visited.add(func_fqn)
+
+    # Find variables that are used but not defined
+    scope = object_scope.get_scope(func_or_var)
+    for item in find_undefined_variables(func_or_var):
+        obj = object_scope.resolve_attribute_to_object(item, scope)
+        info = versions.VersionInfo.of(obj, version_scraping=version_scraping)
+        call_dependencies[info] = obj
+
+        if info.version is None:
+            get_call_dependencies(obj, version_scraping, call_dependencies, visited)
+    return call_dependencies