diff --git a/CHANGES.md b/CHANGES.md index a6b0b443..8b761f98 100644 --- a/CHANGES.md +++ b/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)) diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 963a8259..212e6672 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -749,6 +749,66 @@ 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 annotate is None: + return False + + wrapped = obj_dict.get("__wrapped__") + if wrapped is None: + return False + + try: + wrapped_annotate = wrapped.__annotate__ + except Exception: + return False + + return annotate is wrapped_annotate + + +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 +1462,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 diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index e2097d1c..3bf8393a 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -2707,6 +2707,36 @@ 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__ + + 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