Skip to content

Commit 99aab61

Browse files
authored
gh-104035: Do not ignore user-defined __{get,set}state__ in slotted frozen dataclasses (#104041)
1 parent e147694 commit 99aab61

File tree

3 files changed

+74
-2
lines changed

3 files changed

+74
-2
lines changed

Lib/dataclasses.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -1227,8 +1227,10 @@ def _add_slots(cls, is_frozen, weakref_slot):
12271227

12281228
if is_frozen:
12291229
# Need this for pickling frozen classes with slots.
1230-
cls.__getstate__ = _dataclass_getstate
1231-
cls.__setstate__ = _dataclass_setstate
1230+
if '__getstate__' not in cls_dict:
1231+
cls.__getstate__ = _dataclass_getstate
1232+
if '__setstate__' not in cls_dict:
1233+
cls.__setstate__ = _dataclass_setstate
12321234

12331235
return cls
12341236

Lib/test/test_dataclasses.py

+68
Original file line numberDiff line numberDiff line change
@@ -3184,6 +3184,74 @@ def test_frozen_pickle(self):
31843184
self.assertIsNot(obj, p)
31853185
self.assertEqual(obj, p)
31863186

3187+
@dataclass(frozen=True, slots=True)
3188+
class FrozenSlotsGetStateClass:
3189+
foo: str
3190+
bar: int
3191+
3192+
getstate_called: bool = field(default=False, compare=False)
3193+
3194+
def __getstate__(self):
3195+
object.__setattr__(self, 'getstate_called', True)
3196+
return [self.foo, self.bar]
3197+
3198+
@dataclass(frozen=True, slots=True)
3199+
class FrozenSlotsSetStateClass:
3200+
foo: str
3201+
bar: int
3202+
3203+
setstate_called: bool = field(default=False, compare=False)
3204+
3205+
def __setstate__(self, state):
3206+
object.__setattr__(self, 'setstate_called', True)
3207+
object.__setattr__(self, 'foo', state[0])
3208+
object.__setattr__(self, 'bar', state[1])
3209+
3210+
@dataclass(frozen=True, slots=True)
3211+
class FrozenSlotsAllStateClass:
3212+
foo: str
3213+
bar: int
3214+
3215+
getstate_called: bool = field(default=False, compare=False)
3216+
setstate_called: bool = field(default=False, compare=False)
3217+
3218+
def __getstate__(self):
3219+
object.__setattr__(self, 'getstate_called', True)
3220+
return [self.foo, self.bar]
3221+
3222+
def __setstate__(self, state):
3223+
object.__setattr__(self, 'setstate_called', True)
3224+
object.__setattr__(self, 'foo', state[0])
3225+
object.__setattr__(self, 'bar', state[1])
3226+
3227+
def test_frozen_slots_pickle_custom_state(self):
3228+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
3229+
with self.subTest(proto=proto):
3230+
obj = self.FrozenSlotsGetStateClass('a', 1)
3231+
dumped = pickle.dumps(obj, protocol=proto)
3232+
3233+
self.assertTrue(obj.getstate_called)
3234+
self.assertEqual(obj, pickle.loads(dumped))
3235+
3236+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
3237+
with self.subTest(proto=proto):
3238+
obj = self.FrozenSlotsSetStateClass('a', 1)
3239+
obj2 = pickle.loads(pickle.dumps(obj, protocol=proto))
3240+
3241+
self.assertTrue(obj2.setstate_called)
3242+
self.assertEqual(obj, obj2)
3243+
3244+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
3245+
with self.subTest(proto=proto):
3246+
obj = self.FrozenSlotsAllStateClass('a', 1)
3247+
dumped = pickle.dumps(obj, protocol=proto)
3248+
3249+
self.assertTrue(obj.getstate_called)
3250+
3251+
obj2 = pickle.loads(dumped)
3252+
self.assertTrue(obj2.setstate_called)
3253+
self.assertEqual(obj, obj2)
3254+
31873255
def test_slots_with_default_no_init(self):
31883256
# Originally reported in bpo-44649.
31893257
@dataclass(slots=True)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Do not ignore user-defined ``__getstate__`` and ``__setstate__`` methods for
2+
slotted frozen dataclasses.

0 commit comments

Comments
 (0)