gh-104035: Do not ignore user-defined `__{get,set}state__` in slotted frozen dataclasses (#104041)

This commit is contained in:
Nikita Sobolev 2023-05-01 18:19:06 +03:00 committed by GitHub
parent e147694252
commit 99aab61062
3 changed files with 74 additions and 2 deletions

View File

@ -1227,7 +1227,9 @@ def _add_slots(cls, is_frozen, weakref_slot):
if is_frozen: if is_frozen:
# Need this for pickling frozen classes with slots. # Need this for pickling frozen classes with slots.
if '__getstate__' not in cls_dict:
cls.__getstate__ = _dataclass_getstate cls.__getstate__ = _dataclass_getstate
if '__setstate__' not in cls_dict:
cls.__setstate__ = _dataclass_setstate cls.__setstate__ = _dataclass_setstate
return cls return cls

View File

@ -3184,6 +3184,74 @@ class TestSlots(unittest.TestCase):
self.assertIsNot(obj, p) self.assertIsNot(obj, p)
self.assertEqual(obj, p) self.assertEqual(obj, p)
@dataclass(frozen=True, slots=True)
class FrozenSlotsGetStateClass:
foo: str
bar: int
getstate_called: bool = field(default=False, compare=False)
def __getstate__(self):
object.__setattr__(self, 'getstate_called', True)
return [self.foo, self.bar]
@dataclass(frozen=True, slots=True)
class FrozenSlotsSetStateClass:
foo: str
bar: int
setstate_called: bool = field(default=False, compare=False)
def __setstate__(self, state):
object.__setattr__(self, 'setstate_called', True)
object.__setattr__(self, 'foo', state[0])
object.__setattr__(self, 'bar', state[1])
@dataclass(frozen=True, slots=True)
class FrozenSlotsAllStateClass:
foo: str
bar: int
getstate_called: bool = field(default=False, compare=False)
setstate_called: bool = field(default=False, compare=False)
def __getstate__(self):
object.__setattr__(self, 'getstate_called', True)
return [self.foo, self.bar]
def __setstate__(self, state):
object.__setattr__(self, 'setstate_called', True)
object.__setattr__(self, 'foo', state[0])
object.__setattr__(self, 'bar', state[1])
def test_frozen_slots_pickle_custom_state(self):
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
with self.subTest(proto=proto):
obj = self.FrozenSlotsGetStateClass('a', 1)
dumped = pickle.dumps(obj, protocol=proto)
self.assertTrue(obj.getstate_called)
self.assertEqual(obj, pickle.loads(dumped))
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
with self.subTest(proto=proto):
obj = self.FrozenSlotsSetStateClass('a', 1)
obj2 = pickle.loads(pickle.dumps(obj, protocol=proto))
self.assertTrue(obj2.setstate_called)
self.assertEqual(obj, obj2)
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
with self.subTest(proto=proto):
obj = self.FrozenSlotsAllStateClass('a', 1)
dumped = pickle.dumps(obj, protocol=proto)
self.assertTrue(obj.getstate_called)
obj2 = pickle.loads(dumped)
self.assertTrue(obj2.setstate_called)
self.assertEqual(obj, obj2)
def test_slots_with_default_no_init(self): def test_slots_with_default_no_init(self):
# Originally reported in bpo-44649. # Originally reported in bpo-44649.
@dataclass(slots=True) @dataclass(slots=True)

View File

@ -0,0 +1,2 @@
Do not ignore user-defined ``__getstate__`` and ``__setstate__`` methods for
slotted frozen dataclasses.