diff --git a/Lib/dataclasses.py b/Lib/dataclasses.py index 7a24f8a9e5ccee..c95db22d70c426 100644 --- a/Lib/dataclasses.py +++ b/Lib/dataclasses.py @@ -1184,17 +1184,21 @@ def _process_class(cls, init, repr, eq, order, unsafe_hash, frozen, return cls -# _dataclass_getstate and _dataclass_setstate are needed for pickling frozen -# classes with slots. These could be slightly more performant if we generated +# _dataclass_reduce and _dataclass_produce are needed because pickle will +# crash when there is a cycle in the object graph and the object serves +# as a key for a dict, set or frozenset in one of its descendants. +# cf. gh python/cpython#124937. +# +# These could be slightly more performant if we generated # the code instead of iterating over fields. But that can be a project for # another day, if performance becomes an issue. -def _dataclass_getstate(self): - return [getattr(self, f.name) for f in fields(self)] +def _dataclass_reduce(self): + return _dataclass_produce, ([getattr(self, f.name) for f in fields(self)],) -def _dataclass_setstate(self, state): - for field, value in zip(fields(self), state): - # use setattr because dataclass may be frozen +def _dataclass_produce(self, data): + for field, value in zip(fields(self), data): + # use `object.__setattr__` because dataclass may be frozen. object.__setattr__(self, field.name, value) @@ -1305,12 +1309,8 @@ def _add_slots(cls, is_frozen, weakref_slot, defined_fields): if qualname is not None: newcls.__qualname__ = qualname - if is_frozen: - # Need this for pickling frozen classes with slots. - if '__getstate__' not in cls_dict: - newcls.__getstate__ = _dataclass_getstate - if '__setstate__' not in cls_dict: - newcls.__setstate__ = _dataclass_setstate + if '__reduce__' not in cls_dict: + newcls.__reduce__ = _dataclass_reduce # Fix up any closures which reference __class__. This is used to # fix zero argument super so that it points to the correct class diff --git a/Lib/test/test_dataclasses/__init__.py b/Lib/test/test_dataclasses/__init__.py index 2e6c49e29ce828..4782d419c7c6dd 100644 --- a/Lib/test/test_dataclasses/__init__.py +++ b/Lib/test/test_dataclasses/__init__.py @@ -2194,6 +2194,28 @@ class R: self.assertEqual(new_sample.x, another_new_sample.x) self.assertEqual(sample.y, another_new_sample.y) + def test_dataclasses_hash_pickleable(self): + global P, Q + class Q: + def __init__(self): + self.p = set() + + @dataclass(frozen=True) + class P: + q: Q + + q = Q() + q.p = P(q) + q.p.add(q) + new_q = pickle.loads(pickle.dumps(q)) + + self.assertEqual(q, new_q) + self.assertIsNot(q, new_q) + + self.assertEqual(q.p, new_q.p) + self.assertIsNot(q.p, new_q.p) + + def test_dataclasses_qualnames(self): @dataclass(order=True, unsafe_hash=True, frozen=True) class A: diff --git a/Misc/NEWS.d/next/Library/2024-10-05-10-24-48.gh-issue-125004.noUctj.rst b/Misc/NEWS.d/next/Library/2024-10-05-10-24-48.gh-issue-125004.noUctj.rst new file mode 100644 index 00000000000000..8427d369b9940b --- /dev/null +++ b/Misc/NEWS.d/next/Library/2024-10-05-10-24-48.gh-issue-125004.noUctj.rst @@ -0,0 +1,2 @@ +Fix unpickling for :mod:`dataclasses` with hash-based data structures in their +descendants in the presence of cycles.