Skip to content

Commit

Permalink
pythongh-104035: Do not ignore user-defined __{get,set}state__ in s…
Browse files Browse the repository at this point in the history
…lotted frozen dataclasses (python#104041)
  • Loading branch information
sobolevn authored May 1, 2023
1 parent e147694 commit 99aab61
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 2 deletions.
6 changes: 4 additions & 2 deletions Lib/dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -1227,8 +1227,10 @@ def _add_slots(cls, is_frozen, weakref_slot):

if is_frozen:
# Need this for pickling frozen classes with slots.
cls.__getstate__ = _dataclass_getstate
cls.__setstate__ = _dataclass_setstate
if '__getstate__' not in cls_dict:
cls.__getstate__ = _dataclass_getstate
if '__setstate__' not in cls_dict:
cls.__setstate__ = _dataclass_setstate

return cls

Expand Down
68 changes: 68 additions & 0 deletions Lib/test/test_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -3184,6 +3184,74 @@ def test_frozen_pickle(self):
self.assertIsNot(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):
# Originally reported in bpo-44649.
@dataclass(slots=True)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Do not ignore user-defined ``__getstate__`` and ``__setstate__`` methods for
slotted frozen dataclasses.

0 comments on commit 99aab61

Please sign in to comment.