WIP

Author: crusaderky

src/array_api_extra/_lib/_utils/_helpers.py

@@ -6,7 +6,7 @@
 import math
 import pickle
 from collections.abc import Generator, Iterable
-from types import ModuleType
+from types import ModuleType, NoneType
 from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
 
 from . import _compat
@@ -319,6 +319,12 @@ def capabilities(xp: ModuleType) -> dict[str, int]:
     return out
 
 
+_BASIC_TYPES = frozenset((
+    NoneType, bool, int, float, complex, str, bytes, bytearray,
+    list, tuple, dict, set, frozenset, range, slice,
+))  # fmt: skip
+
+
 def pickle_without(
     obj: object, cls: type[T] | tuple[type[T], ...] = ()
 ) -> tuple[bytes, tuple[T, ...], tuple[object, ...]]:
@@ -390,9 +396,12 @@ class Pickler(pickle.Pickler):  # numpydoc ignore=GL01,RT01
 
         @override
         def persistent_id(self, obj: object) -> object:  # pyright: ignore[reportIncompatibleMethodOverride]  # numpydoc ignore=GL08
+            if type(obj) in _BASIC_TYPES:  # No subclasses!
+                return None
+
             id_ = id(obj)
             try:
-                return seen[id_]
+                return id_, seen[id_]
             except KeyError:
                 pass