Simplify the code

Author: seratch

src/agents/util/_safe_copy.py

@@ -1,104 +1,38 @@
 from __future__ import annotations
 
-import copy
-import datetime as _dt
-from decimal import Decimal
-from fractions import Fraction
-from pathlib import PurePath
 from typing import Any, TypeVar
-from uuid import UUID
 
 T = TypeVar("T")
 
 
 def safe_copy(obj: T) -> T:
     """
-    Copy 'obj' without triggering deepcopy on complex/fragile objects.
-
-    Rules:
-      - Primitive/simple atoms (ints, strs, datetimes, etc.): deepcopy (cheap and safe).
-      - Built-in containers (dict, list, tuple, set, frozenset): recurse element-wise.
-      - Everything else (framework objects, iterators, models, file handles, etc.):
-        shallow copy if possible; otherwise return as-is.
-
+    Craete a copy of the given object -- it can be either str or list/set/tuple of objects.
     This avoids failures like:
       TypeError: cannot pickle '...ValidatorIterator' object
     because we never call deepcopy() on non-trivial objects.
     """
-    memo: dict[int, Any] = {}
-    return _safe_copy_internal(obj, memo)
-
-
-_SIMPLE_ATOMS = (
-    # basics
-    type(None),
-    bool,
-    int,
-    float,
-    complex,
-    str,
-    bytes,
-    # small buffers/scalars
-    bytearray,
-    memoryview,
-    range,
-    # "value" types
-    Decimal,
-    Fraction,
-    UUID,
-    PurePath,
-    _dt.date,
-    _dt.datetime,
-    _dt.time,
-    _dt.timedelta,
-)
-
-
-def _is_simple_atom(o: Any) -> bool:
-    return isinstance(o, _SIMPLE_ATOMS)
-
-
-def _safe_copy_internal(obj: T, memo: dict[int, Any]) -> T:
-    oid = id(obj)
-    if oid in memo:
-        return memo[oid]  # type: ignore [no-any-return]
-
-    # 1) Simple "atoms": safe to deepcopy (cheap, predictable).
-    if _is_simple_atom(obj):
-        return copy.deepcopy(obj)
+    return _safe_copy_internal(obj)
 
-    # 2) Containers: rebuild and recurse.
-    if isinstance(obj, dict):
-        new_dict: dict[Any, Any] = {}
-        memo[oid] = new_dict
-        for k, v in obj.items():
-            # preserve key identity/value, only copy the value
-            new_dict[k] = _safe_copy_internal(v, memo)
-        return new_dict  # type: ignore [return-value]
 
+def _safe_copy_internal(obj: T) -> T:
     if isinstance(obj, list):
         new_list: list[Any] = []
-        memo[oid] = new_list
-        new_list.extend(_safe_copy_internal(x, memo) for x in obj)
+        new_list.extend(_safe_copy_internal(x) for x in obj)
         return new_list  # type: ignore [return-value]
 
     if isinstance(obj, tuple):
-        new_tuple = tuple(_safe_copy_internal(x, memo) for x in obj)
-        memo[oid] = new_tuple
+        new_tuple = tuple(_safe_copy_internal(x) for x in obj)
         return new_tuple  # type: ignore [return-value]
 
     if isinstance(obj, set):
         new_set: set[Any] = set()
-        memo[oid] = new_set
         for x in obj:
-            new_set.add(_safe_copy_internal(x, memo))
+            new_set.add(_safe_copy_internal(x))
         return new_set  # type: ignore [return-value]
 
     if isinstance(obj, frozenset):
-        new_fset = frozenset(_safe_copy_internal(x, memo) for x in obj)
-        memo[oid] = new_fset
+        new_fset = frozenset(_safe_copy_internal(x) for x in obj)
         return new_fset  # type: ignore
 
-    # 3) Unknown/complex leaf: return as-is (identity preserved).
-    memo[oid] = obj
     return obj

tests/utils/test_safe_copy.py

@@ -35,15 +35,6 @@ def __deepcopy__(self, memo):
         raise TypeError("no deepcopy")
 
 
-def test_primitives_are_copied_independently_for_mutable_bytes():
-    orig = bytearray(b"abc")
-    cpy = safe_copy(orig)
-    assert bytes(cpy) == b"abc"
-    orig[0] = ord("z")
-    assert bytes(orig) == b"zbc"
-    assert bytes(cpy) == b"abc"  # unaffected
-
-
 @pytest.mark.parametrize(
     "value",
     [
@@ -69,37 +60,6 @@ def test_simple_atoms_roundtrip(value):
     assert cpy == value
 
 
-def test_deep_copy_for_nested_containers_of_primitives():
-    orig = {"a": [1, 2, {"z": (3, 4)}]}
-    cpy = safe_copy(orig)
-
-    # mutate original deeply
-    orig["a"][2]["z"] = (99, 100)  # type: ignore
-
-    assert cpy == {"a": [1, 2, {"z": (3, 4)}]}  # unaffected
-
-
-def test_complex_leaf_is_only_shallow_copied():
-    class Leaf:
-        def __init__(self):
-            self.val = 1
-
-    leaf = Leaf()
-    obj = {"k": leaf, "arr": [1, 2, 3]}
-    cpy = safe_copy(obj)
-
-    # container structure is new
-    assert cpy is not obj
-    assert cpy["arr"] is not obj["arr"]
-
-    # complex leaf is shallow: identity preserved
-    assert cpy["k"] is leaf
-
-    # mutating the leaf reflects in the copied structure
-    leaf.val = 42
-    assert cpy["k"].val == 42  # type: ignore [attr-defined]
-
-
 def test_generator_is_preserved_and_not_consumed():
     gen = (i for i in range(3))
     data = {"g": gen}
@@ -145,20 +105,6 @@ class Marker:
     assert 99 not in s2
 
 
-def test_cycles_are_handled_without_recursion_error():
-    # a -> (a,)
-    a: list[Any] = []
-    t = (a,)
-    a.append(t)
-
-    c = safe_copy(a)
-    # structure cloned:
-    assert c is not a
-    assert isinstance(c[0], tuple)
-    # cycle preserved: the tuple's first element points back to the list
-    assert c[0][0] is c
-
-
 def test_object_where_deepcopy_would_fail_is_handled_via_shallow_copy():
     b = BoomDeepcopy(7)
     c = safe_copy(b)