WIP

Author: dmontagu

pydantic_graph/pydantic_graph/_utils.py

@@ -3,6 +3,7 @@
 import asyncio
 import types
 import warnings
+import inspect
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from functools import partial
@@ -163,3 +164,41 @@ def logfire_span(*args: Any, **kwargs: Any) -> Generator[LogfireSpan, None, None
             warnings.filterwarnings('ignore', category=LogfireNotConfiguredWarning)
             with _logfire.span(*args, **kwargs) as span:
                 yield span
+
+
+def infer_obj_name(obj: Any, *, depth: int) -> str | None:
+    """Infer the variable name of an object from the calling frame's scope.
+
+    This function examines the call stack to find what variable name was used
+    for the given object in the calling scope. This is useful for automatic
+    naming of objects based on their variable names.
+
+    Args:
+        obj: The object whose variable name to infer.
+        depth: Number of stack frames to traverse upward from the current frame.
+
+    Returns:
+        The inferred variable name if found, None otherwise.
+
+    Example:
+        Usage should generally look like `infer_name(self, depth=2)` or similar.
+    """
+    target_frame = inspect.currentframe()
+    if target_frame is None:
+        return None  # pragma: no cover
+    for _ in range(depth):
+        target_frame = target_frame.f_back
+        if target_frame is None:
+            return None
+
+    for name, item in target_frame.f_locals.items():
+        if item is obj:
+            return name
+
+    if target_frame.f_locals != target_frame.f_globals:  # pragma: no branch
+        # if we couldn't find the agent in locals and globals are a different dict, try globals
+        for name, item in target_frame.f_globals.items():
+            if item is obj:
+                return name
+
+    return None

pydantic_graph/pydantic_graph/beta/graph.py

@@ -22,7 +22,7 @@
 
 from pydantic_ai.exceptions import ExceptionGroup
 from pydantic_graph import exceptions
-from pydantic_graph._utils import AbstractSpan, get_traceparent, logfire_span
+from pydantic_graph._utils import AbstractSpan, get_traceparent, logfire_span, infer_obj_name
 from pydantic_graph.beta.decision import Decision
 from pydantic_graph.beta.id_types import ForkID, ForkStack, ForkStackItem, GraphRunID, JoinID, NodeID, NodeRunID, TaskID
 from pydantic_graph.beta.join import Join, JoinNode, JoinState, ReducerContext
@@ -187,7 +187,9 @@ async def run(
             The final output from the graph execution
         """
         if infer_name and self.name is None:
-            self._infer_name(inspect.currentframe())
+            inferred_name = infer_obj_name(self, depth=2)
+            if inferred_name is not None:
+                self.name = inferred_name
 
         async with self.iter(state=state, deps=deps, inputs=inputs, span=span, infer_name=False) as graph_run:
             # Note: This would probably be better using `async for _ in graph_run`, but this tests the `next` method,
@@ -602,7 +604,7 @@ async def iter_graph(  # noqa C901
                                         for afs in active_fork_stacks
                                     ):
                                         # this join_state is a strict prefix for one of the other active join_states
-                                        continue  # pragma: no cover  # TODO: We should cover this
+                                        continue  # pragma: no cover  # It's difficult to cover this
                                     self.active_reducers.pop(
                                         (join_id, fork_run_id)
                                     )  # we're handling it now, so we can pop it

pydantic_graph/pydantic_graph/beta/util.py

@@ -4,7 +4,6 @@
 including workarounds for type checker limitations and utilities for runtime type inspection.
 """
 
-import inspect
 from dataclasses import dataclass
 from typing import Any, Generic, cast, get_args, get_origin
 
@@ -17,16 +16,16 @@
 class TypeExpression(Generic[T]):
     """A workaround for type checker limitations when using complex type expressions.
 
-    This class serves as a wrapper for types that cannot normally be used in positions
+        This class serves as a wrapper for types that cannot normally be used in positions
     requiring `type[T]`, such as `Any`, `Union[...]`, or `Literal[...]`. It provides a
-    way to pass these complex type expressions to functions expecting concrete types.
+        way to pass these complex type expressions to functions expecting concrete types.
 
-    Example:
-        Instead of `output_type=Union[str, int]` (which may cause type errors),
-        use `output_type=TypeExpression[Union[str, int]]`.
+        Example:
+            Instead of `output_type=Union[str, int]` (which may cause type errors),
+            use `output_type=TypeExpression[Union[str, int]]`.
 
-    Note:
-        This is a workaround for the lack of TypeForm in the Python type system.
+        Note:
+            This is a workaround for the lack of TypeForm in the Python type system.
     """
 
     pass
@@ -89,44 +88,3 @@ def get_callable_name(callable_: Any) -> str:
         The callable's __name__ attribute if available, otherwise its string representation.
     """
     return getattr(callable_, '__name__', str(callable_))
-
-
-def infer_name(obj: Any, *, depth: int) -> str | None:
-    """Infer the variable name of an object from the calling frame's scope.
-
-    This function examines the call stack to find what variable name was used
-    for the given object in the calling scope. This is useful for automatic
-    naming of objects based on their variable names.
-
-    Args:
-        obj: The object whose variable name to infer.
-        depth: Number of stack frames to traverse upward from the current frame.
-
-    Returns:
-        The inferred variable name if found, None otherwise.
-
-    Example:
-        Usage should generally look like `infer_name(self, depth=2)` or similar.
-
-    Note:
-        TODO(P3): Use this or lose it
-    """
-    target_frame = inspect.currentframe()
-    if target_frame is None:
-        return None  # pragma: no cover
-    for _ in range(depth):
-        target_frame = target_frame.f_back
-        if target_frame is None:
-            return None
-
-    for name, item in target_frame.f_locals.items():
-        if item is obj:
-            return name
-
-    if target_frame.f_locals != target_frame.f_globals:  # pragma: no branch
-        # if we couldn't find the agent in locals and globals are a different dict, try globals
-        for name, item in target_frame.f_globals.items():
-            if item is obj:
-                return name
-
-    return None

tests/graph/beta/test_util.py

@@ -4,7 +4,6 @@
     Some,
     TypeExpression,
     get_callable_name,
-    infer_name,
     unpack_type_expression,
 )
 
@@ -48,45 +47,3 @@ class MyClass:
     name = get_callable_name(obj)
     assert isinstance(name, str)
     assert 'object' in name
-
-
-def test_infer_name():
-    """Test inferring variable names from the calling frame."""
-    my_object = object()
-    # Depth 1 means we look at the frame calling infer_name
-    inferred = infer_name(my_object, depth=1)
-    assert inferred == 'my_object'
-
-    # Test with object not in locals
-    result = infer_name(object(), depth=1)
-    assert result is None
-
-
-def test_infer_name_no_frame():
-    """Test infer_name when frame inspection fails."""
-    # This is hard to trigger without mocking, but we can test that the function
-    # returns None gracefully when it can't find the object
-    some_obj = object()
-
-    # Call with depth that would exceed the call stack
-    result = infer_name(some_obj, depth=1000)
-    assert result is None
-
-
-global_obj = object()
-
-
-def test_infer_name_locals_vs_globals():
-    """Test infer_name prefers locals over globals."""
-    result = infer_name(global_obj, depth=1)
-    assert result == 'global_obj'
-
-    # Assign a local name to the variable and ensure it is found with precedence over the global
-    local_obj = global_obj
-    result = infer_name(global_obj, depth=1)
-    assert result == 'local_obj'
-
-    # If we unbind the local name, should find the global name again
-    del local_obj
-    result = infer_name(global_obj, depth=1)
-    assert result == 'global_obj'

tests/graph/test_utils.py

@@ -1,6 +1,6 @@
 from threading import Thread
 
-from pydantic_graph._utils import get_event_loop
+from pydantic_graph._utils import get_event_loop, infer_obj_name
 
 
 def test_get_event_loop_in_thread():
@@ -11,3 +11,45 @@ def get_and_close_event_loop():
     thread = Thread(target=get_and_close_event_loop)
     thread.start()
     thread.join()
+
+
+def test_infer_obj_name():
+    """Test inferring variable names from the calling frame."""
+    my_object = object()
+    # Depth 1 means we look at the frame calling infer_obj_name
+    inferred = infer_obj_name(my_object, depth=1)
+    assert inferred == 'my_object'
+
+    # Test with object not in locals
+    result = infer_obj_name(object(), depth=1)
+    assert result is None
+
+
+def test_infer_obj_name_no_frame():
+    """Test infer_obj_name when frame inspection fails."""
+    # This is hard to trigger without mocking, but we can test that the function
+    # returns None gracefully when it can't find the object
+    some_obj = object()
+
+    # Call with depth that would exceed the call stack
+    result = infer_obj_name(some_obj, depth=1000)
+    assert result is None
+
+
+global_obj = object()
+
+
+def test_infer_obj_name_locals_vs_globals():
+    """Test infer_obj_name prefers locals over globals."""
+    result = infer_obj_name(global_obj, depth=1)
+    assert result == 'global_obj'
+
+    # Assign a local name to the variable and ensure it is found with precedence over the global
+    local_obj = global_obj
+    result = infer_obj_name(global_obj, depth=1)
+    assert result == 'local_obj'
+
+    # If we unbind the local name, should find the global name again
+    del local_obj
+    result = infer_obj_name(global_obj, depth=1)
+    assert result == 'global_obj'