Allow pickling Cancelled

newsfragments/3248.bugfix.rst

@@ -0,0 +1 @@
+Allow pickling `trio.Cancelled`, as they can show up when you want to pickle something else. Also enables pickling for other ``NoPublicConstructor`` objects.

src/trio/_core/_exceptions.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from trio._util import NoPublicConstructor, final
 
 

src/trio/_core/_tests/test_run.py

@@ -3,6 +3,7 @@
 import contextvars
 import functools
 import gc
+import pickle
 import sys
 import threading
 import time
@@ -2235,6 +2236,13 @@ def test_Cancelled_subclass() -> None:
         type("Subclass", (_core.Cancelled,), {})
 
 
+# https://github.com/python-trio/trio/issues/3248
+def test_Cancelled_pickle() -> None:
+    cancelled = _core.Cancelled._create()
+    cancelled = pickle.loads(pickle.dumps(cancelled))
+    assert isinstance(cancelled, _core.Cancelled)
+
+
 def test_CancelScope_subclass() -> None:
     with pytest.raises(TypeError):
         type("Subclass", (_core.CancelScope,), {})

src/trio/_tests/test_util.py

@@ -2,7 +2,7 @@
 
 import sys
 import types
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Coroutine, Generator
@@ -198,11 +198,16 @@ def __init__(self, a: int, b: float) -> None:
             assert a == 8
             assert b == 3.15
 
+        def __reduce__(self) -> tuple[Any, ...]:  # type: ignore[explicit-any]
+            return ("dummy", ())
+
     with pytest.raises(TypeError):
         SpecialClass(8, 3.15)
 
     # Private constructor should not raise, and passes args to __init__.
-    assert isinstance(SpecialClass._create(8, b=3.15), SpecialClass)
+    inst = SpecialClass._create(8, b=3.15)
+    assert isinstance(inst, SpecialClass)
+    assert inst.__reduce__()[0] == SpecialClass._create
 
 
 def test_fixup_module_metadata() -> None:

src/trio/_util.py

@@ -327,6 +327,21 @@ def __call__(cls, *args: object, **kwargs: object) -> None:
     def _create(cls: type[T], *args: object, **kwargs: object) -> T:
         return super().__call__(*args, **kwargs)  # type: ignore
 
+    def __new__(cls, name: str, bases: tuple[type, ...], namespace: dict[str, Any], **kwargs: object) -> type:  # type: ignore[explicit-any]
+        old_reduce = namespace.get("__reduce__")
+
+        def __reduce__(self: object) -> str | tuple[Any, ...]:  # type: ignore[explicit-any]
+            if old_reduce is not None:
+                result = old_reduce(self)
+                assert not isinstance(result, str)
+            else:
+                result = super(type(self), self).__reduce__()  # type: ignore[misc]
+
+            return (type(self)._create, *result[1:])  # type: ignore[attr-defined]
+
+        namespace["__reduce__"] = __reduce__
+        return super().__new__(cls, name, bases, namespace, **kwargs)
+
 
 def name_asyncgen(agen: AsyncGeneratorType[object, NoReturn]) -> str:
     """Return the fully-qualified name of the async generator function