Fix update_wrapper pickling with lazy annotations

Author: BioGeek

CHANGES.md

@@ -1,6 +1,10 @@
 In development
 ==============
 
+- Fix pickling of objects decorated with `functools.update_wrapper` when the
+  wrapped callable uses Python 3.14 lazy annotations. ([issue #585](
+  https://github.com/cloudpipe/cloudpickle/issues/585))
+
 - Make pickling of functions depending on globals in notebook more
   deterministic. ([PR#560](https://github.com/cloudpipe/cloudpickle/pull/560))
 

cloudpickle/cloudpickle.py

@@ -749,6 +749,73 @@ def _function_getstate(func):
     return state, slotstate
 
 
+def _is_copied_annotation_function(obj):
+    """Detect __annotate__ copied by functools.update_wrapper on Python 3.14+."""
+    if sys.version_info < (3, 14):
+        return False
+
+    try:
+        obj_dict = obj.__dict__
+    except Exception:
+        return False
+    if not isinstance(obj_dict, dict):
+        return False
+
+    annotate = obj_dict.get("__annotate__")
+    if not isinstance(annotate, types.FunctionType):
+        return False
+
+    wrapped = obj_dict.get("__wrapped__")
+    if wrapped is None:
+        return False
+
+    if annotate.__name__ != "__annotate__":
+        return False
+
+    annotate_module = getattr(annotate, "__module__", None)
+    wrapped_module = getattr(wrapped, "__module__", None)
+    if (
+        annotate_module is not None
+        and wrapped_module is not None
+        and annotate_module != wrapped_module
+    ):
+        return False
+
+    return True
+
+
+def _remove_key_from_state(state, key):
+    if isinstance(state, dict):
+        if key not in state:
+            return state
+        state = state.copy()
+        state.pop(key, None)
+        return state
+
+    if (
+        isinstance(state, tuple)
+        and len(state) == 2
+        and isinstance(state[0], dict)
+        and key in state[0]
+    ):
+        state_dict = state[0].copy()
+        state_dict.pop(key, None)
+        return state_dict, state[1]
+
+    return state
+
+
+def _reduced_without_copied_annotation_function(obj, proto):
+    """Remove redundant __annotate__ copied from a wrapped callable."""
+    rv = obj.__reduce_ex__(proto)
+    if not isinstance(rv, tuple) or len(rv) < 3:
+        return rv
+
+    state = rv[2]
+    state = _remove_key_from_state(state, "__annotate__")
+    return rv[:2] + (state,) + rv[3:]
+
+
 def _class_getstate(obj):
     clsdict = _extract_class_dict(obj)
     clsdict.pop("__weakref__", None)
@@ -1402,6 +1469,8 @@ def reducer_override(self, obj):
                 return _class_reduce(obj)
             elif isinstance(obj, types.FunctionType):
                 return self._function_reduce(obj)
+            elif _is_copied_annotation_function(obj):
+                return _reduced_without_copied_annotation_function(obj, self.proto)
             else:
                 # fallback to save_global, including the Pickler's
                 # dispatch_table

tests/cloudpickle_test.py

@@ -2707,6 +2707,51 @@ class C(abc.ABC):
         c2 = C2()
         assert isinstance(c2, C2)
 
+    @pytest.mark.skipif(
+        sys.version_info < (3, 14),
+        reason="functools.update_wrapper copies __annotate__ starting in Python 3.14",
+    )
+    def test_update_wrapper_with_annotated_abc_method(self):
+        # see https://github.com/cloudpipe/cloudpickle/issues/585
+        class FuncWrapper:
+            def __init__(self, function):
+                self.function = function
+                functools.update_wrapper(self, self.function)
+
+            def __call__(self, *args, **kwargs):
+                return self.function(*args, **kwargs)
+
+        class AbstractClass(abc.ABC):
+            a: int
+
+            def method(self, arg: str) -> str:
+                return arg.upper()
+
+        wrapped = FuncWrapper(AbstractClass().method)
+        assert "__annotate__" in wrapped.__dict__
+
+        # Simulate CPython builds where the annotation function copied by
+        # update_wrapper is distinct from __wrapped__.__annotate__.
+        copied_annotate = wrapped.__dict__["__annotate__"]
+        wrapped.__annotate__ = types.FunctionType(
+            copied_annotate.__code__,
+            copied_annotate.__globals__,
+            copied_annotate.__name__,
+            copied_annotate.__defaults__,
+            copied_annotate.__closure__,
+        )
+        wrapped.__annotate__.__kwdefaults__ = copied_annotate.__kwdefaults__
+        wrapped.__annotate__.__qualname__ = copied_annotate.__qualname__
+        wrapped.__annotate__.__module__ = copied_annotate.__module__
+        assert wrapped.__annotate__ is not wrapped.__wrapped__.__annotate__
+
+        wrapped_clone = pickle_depickle(wrapped, protocol=self.protocol)
+
+        assert wrapped_clone("abc") == "ABC"
+        assert wrapped_clone.__name__ == "method"
+        assert "__annotate__" not in wrapped_clone.__dict__
+        assert wrapped_clone.__wrapped__.__annotations__ == {"arg": str, "return": str}
+
     def test_function_annotations(self):
         def f(a: int) -> str:
             pass