support unpickleable

Author: crusaderky

src/array_api_extra/_lib/_utils/_helpers.py

@@ -2,14 +2,12 @@
 
 from __future__ import annotations
 
-import copyreg
 import io
 import math
 import pickle
-from collections.abc import Callable, Generator, Iterable, Iterator
-from contextvars import ContextVar
+from collections.abc import Generator, Iterable
 from types import ModuleType
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
 
 from . import _compat
 from ._compat import (
@@ -23,8 +21,13 @@
 from ._typing import Array
 
 if TYPE_CHECKING:  # pragma: no cover
-    # TODO import from typing (requires Python >=3.13)
-    from typing_extensions import TypeIs
+    # TODO import from typing (requires Python >=3.12 and >=3.13)
+    from typing_extensions import TypeIs, override
+else:
+
+    def override(func):
+        return func
+
 
 T = TypeVar("T")
 
@@ -316,48 +319,33 @@ def capabilities(xp: ModuleType) -> dict[str, int]:
     return out
 
 
-# Helper of ``extract_objects`` and ``repack_objects``
-_repacking_objects: ContextVar[Iterator[object]] = ContextVar("_repacking_objects")
-
-
-def _expand() -> object:  # numpydoc ignore=RT01
-    """
-    Helper of ``extract_objects`` and ``repack_objects``.
-
-    Inverse of the reducer function.
-
-    Notes
-    -----
-    This function must be global in order to be picklable.
-    """
-    try:
-        return next(_repacking_objects.get())
-    except StopIteration:
-        msg = "Not enough objects to repack"
-        raise ValueError(msg)
-
-
-def pickle_without(obj: object, *classes: type[T]) -> tuple[bytes, list[T]]:
+def pickle_without(
+    obj: object, cls: type[T] | tuple[type[T], ...] = ()
+) -> tuple[bytes, tuple[T, ...], tuple[object, ...]]:
     """
-    Variant of ``pickle.dumps`` that extracts inner objects.
+    Variant of ``pickle.dumps`` that always succeeds and extracts inner objects.
 
     Conceptually, this is similar to passing the ``buffer_callback`` argument to
-    ``pickle.dumps``, but instead of extracting buffers it extracts entire objects.
+    ``pickle.dumps``, but instead of extracting buffers it extracts entire objects,
+    which are either not serializable with ``pickle`` (e.g. local classes or functions)
+    or instances of an explicit list of classes.
 
     Parameters
     ----------
     obj : object
         The object to pickle.
-    *classes : type
-        One or more classes to extract from the object.
+    cls : type | tuple[type, ...], optional
+        One or multiple classes to extract from the object.
         The instances of these classes inside ``obj`` will not be pickled.
 
     Returns
     -------
     bytes
         The pickled object. Must be unpickled with :func:`unpickle_without`.
-    list
-        All instances of ``classes`` found inside ``obj`` (not pickled).
+    tuple
+        All instances of ``cls`` found inside ``obj`` (not pickled).
+    tuple
+        All other objects which failed to pickle.
 
     See Also
     --------
@@ -366,75 +354,144 @@ def pickle_without(obj: object, *classes: type[T]) -> tuple[bytes, list[T]]:
 
     Examples
     --------
+    >>> class NS:
+    ...     def __repr__(self):
+    ...         return "<NS>"
+    ...     def __reduce__(self):
+    ...         assert False
     >>> class A:
     ...     def __repr__(self):
     ...         return "<A>"
-    ...     def __reduce__(self):
-    ...         assert False, "Not serializable"
-    >>> obj = {1: A(), 2: [A(), A()]}  # Any serializable object
-    >>> pik, extracted = pickle_without(obj, A)
-    >>> extracted
-    [<A>, <A>, <A>]
-    >>> unpickle_without(pik, extracted)
-    {1: <A>, 2: [<A>, <A>]}
+    >>> obj = {1: A(), 2: [A(), NS(), A()]}  # Any serializable object
+    >>> pik, instances, unpickleable = pickle_without(obj, A)
+    >>> instances, unpickleable
+    ([<A>, <A>, <A>], [<NS>])
+    >>> unpickle_without(pik, instances, unpickleable)
+    {1: <A>, 2: [<A>, <NS>, <A>]}
 
     This can be also used to hot-swap inner objects; the only constraint is that
     the number of objects in and out must be the same:
 
     >>> class B:
     ...     def __repr__(self): return "<B>"
-    >>> unpickle_without(pik, [B(), B(), B()])
-    {1: <B>, 2: [<B>, <B>]}
+    >>> unpickle_without(pik, [B(), B(), B()], [NS()])
+    {1: <B>, 2: [<B>, <NS>, <B>]}
     """
-    extracted = []
-
-    def reduce(x: T) -> tuple[Callable[[], object], tuple[()]]:  # numpydoc ignore=GL08
-        extracted.append(x)
-        return _expand, ()
+    instances: list[T] = []
+    unpickleable: list[object] = []
+    seen: dict[int, Literal[0, 1, None]] = {}
+
+    class Pickler(pickle.Pickler):  # numpydoc ignore=GL01,RT01
+        """Override pickle.Pickler.persistent_id.
+
+        TODO consider moving to top-level scope to allow for
+        the full Pickler API to be used.
+        """
+
+        @override
+        def persistent_id(self, obj: object) -> object:  # pyright: ignore[reportIncompatibleMethodOverride]  # numpydoc ignore=GL08
+            id_ = id(obj)
+            try:
+                return seen[id_]
+            except KeyError:
+                pass
+
+            if isinstance(obj, cls):
+                instances.append(obj)  # type: ignore[arg-type]
+                seen[id_] = 0
+                return id_, 0
+
+            try:
+                _ = obj.__reduce__()
+            except Exception:  # pylint: disable=broad-exception-caught
+                pass
+            else:  # Can be pickled
+                seen[id_] = None
+                return None
+
+            # May be a global function, which can be pickled
+            try:
+                _ = pickle.dumps(obj)
+            except Exception:  # pylint: disable=broad-exception-caught
+                pass
+            else:  # Can be pickled
+                seen[id_] = None
+                return None
+
+            # Can't be pickled
+            unpickleable.append(obj)
+            seen[id_] = 1
+            return id_, 1
 
     f = io.BytesIO()
-    p = pickle.Pickler(f)
-
-    # Override the reducer for the given classes and all their
-    # subclasses (recursively).
-    p.dispatch_table = copyreg.dispatch_table.copy()
-    subclasses = list(classes)
-    while subclasses:
-        cls = subclasses.pop()
-        p.dispatch_table[cls] = reduce
-        subclasses.extend(cls.__subclasses__())
-
+    p = Pickler(f)
     p.dump(obj)
+    return f.getvalue(), tuple(instances), tuple(unpickleable)
 
-    return f.getvalue(), extracted
 
-
-def unpickle_without(pik: bytes, objects: Iterable[object], /) -> Any:  # type: ignore[explicit-any]
+def unpickle_without(  # type: ignore[explicit-any]
+    pik: bytes,
+    instances: Iterable[object],
+    unpickleable: Iterable[object],
+    /,
+) -> Any:
     """
     Variant of ``pickle.loads``, reverse of ``pickle_without``.
 
     Parameters
     ----------
     pik : bytes
         The pickled object generated by ``pickle_without``.
-    objects : Iterable
-        The objects to be reinserted into the unpickled object.
-        Must be the at least the same number of elements as the ones extracted by
-        ``pickle_without``, but does not need to be the same objects or even the
-        same types of objects. Excess objects, if any, won't be inserted.
+    instances : Iterable[object]
+        Instances of the class or classes explicitly passed to ``pickle_without``,
+        to be reinserted into the unpickled object.
+    unpickleable : Iterable[object]
+        The objects that failed to pickle, as returned by ``pickle_without``.
 
     Returns
     -------
     object
-        The unpickled object, with the objects in ``objects`` inserted back into it.
+        The unpickled object.
 
     See Also
     --------
     pickle_without : Serializing function.
     pickle.loads : Standard unpickle function.
+
+    Notes
+    -----
+    The second and third parameter of this function must yield at least the same number
+    of elements as the ones returned by ``pickle_without``, but do not need to be the
+    same objects, or even the same types of objects. Excess objects, if any, will be
+    quietly ignored.
     """
-    tok = _repacking_objects.set(iter(objects))
-    try:
-        return pickle.loads(pik)
-    finally:
-        _repacking_objects.reset(tok)
+    iters = iter(instances), iter(unpickleable)
+    seen: dict[int, object] = {}
+
+    class Unpickler(pickle.Unpickler):  # numpydoc ignore=GL01,RT01
+        """
+        Override pickle.Pickler.persistent_load.
+
+        TODO consider moving to top-level scope to allow for
+        the full Unpickler API to be used.
+        """
+
+        @override
+        def persistent_load(self, pid: tuple[int, int]) -> object:  # pyright: ignore[reportIncompatibleMethodOverride]  # numpydoc ignore=GL08
+            prev_id, kind = pid
+            try:
+                return seen[prev_id]
+            except KeyError:
+                pass
+
+            try:
+                obj = next(iters[kind])
+            except StopIteration as e:
+                msg = "Not enough objects to unpickle"
+                raise ValueError(msg) from e
+
+            seen[prev_id] = obj
+            return obj
+
+    f = io.BytesIO(pik)
+    return Unpickler(f).load()

src/array_api_extra/testing.py

@@ -329,9 +329,9 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
         # Block until the graph materializes and reraise exceptions. This allows
         # `pytest.raises` and `pytest.warns` to work as expected. Note that this would
         # not work on scheduler='distributed', as it would not block.
-        pik, arrays = pickle_without(out, da.Array)
+        pik, arrays, unpickleable = pickle_without(out, da.Array)
         arrays = dask.persist(arrays, scheduler="threads")[0]  # type: ignore[attr-defined,no-untyped-call,func-returns-value,index]  # pyright: ignore[reportPrivateImportUsage]
-        return unpickle_without(pik, arrays)
+        return unpickle_without(pik, arrays, unpickleable)  # pyright: ignore[reportUnknownArgumentType]
 
     return wrapper
 
@@ -343,6 +343,8 @@ def _jax_autojit(func: Callable[P, T]) -> Callable[P, T]:  # numpydoc ignore=PR0
     - Array-like arguments and return values are not automatically converted to
       ``jax.Array`` objects.
     - All non-array arguments are automatically treated as static.
+      Unlike ``jax.jit``, static arguments must be either hashable or serializable with
+      ``pickle``.
     - Unlike ``jax.jit``, non-array arguments and return values are not limited to
       tuple/list/dict, but can be any object serializable with ``pickle``.
     - Automatically descend into non-array arguments and find ``jax.Array`` objects
@@ -354,27 +356,44 @@ def _jax_autojit(func: Callable[P, T]) -> Callable[P, T]:  # numpydoc ignore=PR0
     import jax
 
     # pickled return values of `func`, minus the JAX arrays
-    res_piks = {}
-
-    def jit_cache_key(args_pik: bytes, *arg_arrays: jax.Array) -> Hashable:  # type: ignore[no-any-unimported]
-        return (args_pik, *((arr.shape, arr.dtype) for arr in arg_arrays))
-
-    def inner(args_pik: bytes, *arg_arrays: jax.Array) -> list[jax.Array]:  # type: ignore[no-any-unimported]  # numpydoc ignore=GL08
-        args, kwargs = unpickle_without(args_pik, arg_arrays)
-        res = func(*args, **kwargs)
-        res_pik, res_arrays = pickle_without(res, jax.Array)
-        key = jit_cache_key(args_pik, *arg_arrays)
-        prev = res_piks.setdefault(key, res_pik)
-        assert prev == res_pik, "cache key collision"
+    static_return_values: dict[Hashable, tuple[bytes, tuple[object, ...]]] = {}
+
+    def jit_cache_key(  # type: ignore[no-any-unimported]
+        args_pik: bytes,
+        args_arrays: tuple[jax.Array, ...],  # pyright: ignore[reportUnknownParameterType]
+        args_unpickleable: tuple[Hashable, ...],
+    ) -> Hashable:
+        return (
+            args_pik,
+            tuple((arr.shape, arr.dtype) for arr in args_arrays),  # pyright: ignore[reportUnknownArgumentType]
+            args_unpickleable,
+        )
+
+    def inner(  # type: ignore[no-any-unimported]  # pyright: ignore[reportUnknownParameterType]
+        args_pik: bytes,
+        args_arrays: tuple[jax.Array, ...],  # pyright: ignore[reportUnknownParameterType]
+        args_unpickleable: tuple[Hashable, ...],
+    ) -> tuple[jax.Array, ...]:  # numpydoc ignore=GL08
+        args, kwargs = unpickle_without(args_pik, args_arrays, args_unpickleable)  # pyright: ignore[reportUnknownArgumentType]
+        res = func(*args, **kwargs)  # pyright: ignore[reportCallIssue]
+        res_pik, res_arrays, res_unpickleable = pickle_without(res, jax.Array)  # pyright: ignore[reportUnknownArgumentType]
+        key = jit_cache_key(args_pik, args_arrays, args_unpickleable)
+        val = res_pik, res_unpickleable
+        prev = static_return_values.setdefault(key, val)
+        assert prev == val, "cache key collision"
         return res_arrays
 
-    jitted = jax.jit(inner, static_argnums=0)
+    jitted = jax.jit(inner, static_argnums=(0, 2))
 
     @wraps(func)
     def outer(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
-        args_pik, arg_arrays = pickle_without((args, kwargs), jax.Array)
-        res_arrays = jitted(args_pik, *arg_arrays)
-        res_pik = res_piks[jit_cache_key(args_pik, *arg_arrays)]
-        return unpickle_without(res_pik, res_arrays)
+        args_pik, args_arrays, args_unpickleable = pickle_without(
+            (args, kwargs),
+            jax.Array,  # pyright: ignore[reportUnknownArgumentType]
+        )
+        res_arrays = jitted(args_pik, args_arrays, args_unpickleable)
+        key = jit_cache_key(args_pik, args_arrays, args_unpickleable)
+        res_pik, res_unpickleable = static_return_values[key]
+        return unpickle_without(res_pik, res_arrays, res_unpickleable)  # pyright: ignore[reportUnknownArgumentType]
 
     return outer

tests/test_testing.py

@@ -194,9 +194,9 @@ def test_lazy_xp_function_static_params(xp: ModuleType):
 
 def test_lazy_xp_function_deprecated_static_argnames():
     with pytest.warns(DeprecationWarning, match="static_argnames"):
-        lazy_xp_function(static_params, static_argnames=["flag"])  # type: ignore[arg-type]
+        lazy_xp_function(static_params, static_argnames=["flag"])  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
     with pytest.warns(DeprecationWarning, match="static_argnums"):
-        lazy_xp_function(static_params, static_argnums=[1])  # type: ignore[arg-type]
+        lazy_xp_function(static_params, static_argnums=[1])  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
 
 try: