Remove _utils.Either (#768)

Author: dmontagu

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -8,7 +8,7 @@
 from datetime import datetime, timezone
 from functools import partial
 from types import GenericAlias
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union
 
 from pydantic import BaseModel
 from pydantic.json_schema import JsonSchemaValue
@@ -79,61 +79,6 @@ def is_set(t_or_unset: T | Unset) -> TypeGuard[T]:
     return t_or_unset is not UNSET
 
 
-Left = TypeVar('Left')
-Right = TypeVar('Right')
-
-
-class Either(Generic[Left, Right]):
-    """Two member Union that records which member was set, this is analogous to Rust enums with two variants.
-
-    Usage:
-
-    ```python
-    if left_thing := either.left:
-        use_left(left_thing.value)
-    else:
-        use_right(either.right)
-    ```
-    """
-
-    __slots__ = '_left', '_right'
-
-    @overload
-    def __init__(self, *, left: Left) -> None: ...
-
-    @overload
-    def __init__(self, *, right: Right) -> None: ...
-
-    def __init__(self, left: Left | Unset = UNSET, right: Right | Unset = UNSET) -> None:
-        if left is not UNSET:
-            assert right is UNSET, '`Either` must receive exactly one argument - `left` or `right`'
-            self._left: Option[Left] = Some(cast(Left, left))
-        else:
-            assert right is not UNSET, '`Either` must receive exactly one argument - `left` or `right`'
-            self._left = None
-            self._right = cast(Right, right)
-
-    @property
-    def left(self) -> Option[Left]:
-        return self._left
-
-    @property
-    def right(self) -> Right:
-        return self._right
-
-    def is_left(self) -> bool:
-        return self._left is not None
-
-    def whichever(self) -> Left | Right:
-        return self._left.value if self._left is not None else self.right
-
-    def __repr__(self):
-        if left := self._left:
-            return f'Either(left={left.value!r})'
-        else:
-            return f'Either(right={self.right!r})'
-
-
 @asynccontextmanager
 async def group_by_temporal(
     aiterable: AsyncIterable[T], soft_max_interval: float | None

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -34,6 +34,20 @@
 from .function import _estimate_string_tokens, _estimate_usage  # pyright: ignore[reportPrivateUsage]
 
 
+@dataclass
+class _TextResult:
+    """A private wrapper class to tag a result that came from the custom_result_text field."""
+
+    value: str | None
+
+
+@dataclass
+class _FunctionToolResult:
+    """A wrapper class to tag a result that came from the custom_result_args field."""
+
+    value: Any | None
+
+
 @dataclass
 class TestModel(Model):
     """A model specifically for testing purposes.
@@ -53,7 +67,7 @@ class TestModel(Model):
     call_tools: list[str] | Literal['all'] = 'all'
     """List of tools to call. If `'all'`, all tools will be called."""
     custom_result_text: str | None = None
-    """If set, this text is return as the final result."""
+    """If set, this text is returned as the final result."""
     custom_result_args: Any | None = None
     """If set, these args will be passed to the result tool."""
     seed: int = 0
@@ -95,21 +109,21 @@ async def agent_model(
         if self.custom_result_text is not None:
             assert allow_text_result, 'Plain response not allowed, but `custom_result_text` is set.'
             assert self.custom_result_args is None, 'Cannot set both `custom_result_text` and `custom_result_args`.'
-            result: _utils.Either[str | None, Any | None] = _utils.Either(left=self.custom_result_text)
+            result: _TextResult | _FunctionToolResult = _TextResult(self.custom_result_text)
         elif self.custom_result_args is not None:
             assert result_tools is not None, 'No result tools provided, but `custom_result_args` is set.'
             result_tool = result_tools[0]
 
             if k := result_tool.outer_typed_dict_key:
-                result = _utils.Either(right={k: self.custom_result_args})
+                result = _FunctionToolResult({k: self.custom_result_args})
             else:
-                result = _utils.Either(right=self.custom_result_args)
+                result = _FunctionToolResult(self.custom_result_args)
         elif allow_text_result:
-            result = _utils.Either(left=None)
+            result = _TextResult(None)
         elif result_tools:
-            result = _utils.Either(right=None)
+            result = _FunctionToolResult(None)
         else:
-            result = _utils.Either(left=None)
+            result = _TextResult(None)
 
         return TestAgentModel(tool_calls, result, result_tools, self.seed)
 
@@ -126,7 +140,7 @@ class TestAgentModel(AgentModel):
 
     tool_calls: list[tuple[str, ToolDefinition]]
     # left means the text is plain text; right means it's a function call
-    result: _utils.Either[str | None, Any | None]
+    result: _TextResult | _FunctionToolResult
     result_tools: list[ToolDefinition]
     seed: int
     model_name: str = 'test'
@@ -176,16 +190,18 @@ def _request(self, messages: list[ModelMessage], model_settings: ModelSettings |
                         [
                             ToolCallPart.from_raw_args(
                                 tool.name,
-                                self.result.right if self.result.right is not None else self.gen_tool_args(tool),
+                                self.result.value
+                                if isinstance(self.result, _FunctionToolResult) and self.result.value is not None
+                                else self.gen_tool_args(tool),
                             )
                             for tool in self.result_tools
                             if tool.name in new_retry_names
                         ]
                     )
                 return ModelResponse(parts=retry_parts, model_name=self.model_name)
 
-        if response_text := self.result.left:
-            if response_text.value is None:
+        if isinstance(self.result, _TextResult):
+            if (response_text := self.result.value) is None:
                 # build up details of tool responses
                 output: dict[str, Any] = {}
                 for message in messages:
@@ -200,10 +216,10 @@ def _request(self, messages: list[ModelMessage], model_settings: ModelSettings |
                 else:
                     return ModelResponse(parts=[TextPart('success (no tool calls)')], model_name=self.model_name)
             else:
-                return ModelResponse(parts=[TextPart(response_text.value)], model_name=self.model_name)
+                return ModelResponse(parts=[TextPart(response_text)], model_name=self.model_name)
         else:
             assert self.result_tools, 'No result tools provided'
-            custom_result_args = self.result.right
+            custom_result_args = self.result.value
             result_tool = self.result_tools[self.seed % len(self.result_tools)]
             if custom_result_args is not None:
                 return ModelResponse(

tests/test_utils.py

@@ -9,7 +9,7 @@
 from inline_snapshot import snapshot
 
 from pydantic_ai import UserError
-from pydantic_ai._utils import UNSET, Either, PeekableAsyncStream, check_object_json_schema, group_by_temporal
+from pydantic_ai._utils import UNSET, PeekableAsyncStream, check_object_json_schema, group_by_temporal
 
 from .models.mock_async_stream import MockAsyncStream
 
@@ -49,13 +49,6 @@ def test_check_object_json_schema():
         check_object_json_schema(array_schema)
 
 
-def test_either():
-    assert repr(Either[int, int](left=123)) == 'Either(left=123)'
-    assert Either(left=123).whichever() == 123
-    assert repr(Either[int, int](right=456)) == 'Either(right=456)'
-    assert Either(right=456).whichever() == 456
-
-
 @pytest.mark.parametrize('peek_first', [True, False])
 @pytest.mark.anyio
 async def test_peekable_async_stream(peek_first: bool):