Skip to content

Commit 9141fed

Browse files
sobolevnmiss-islington
authored andcommitted
pythongh-104035: Do not ignore user-defined __{get,set}state__ in slotted frozen dataclasses (pythonGH-104041)
(cherry picked from commit 99aab61) Co-authored-by: Nikita Sobolev <[email protected]>
1 parent 1be8bed commit 9141fed

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
@@ -1191,8 +1191,10 @@ def _add_slots(cls, is_frozen, weakref_slot):
11911191

11921192
if is_frozen:
11931193
# Need this for pickling frozen classes with slots.
1194-
cls.__getstate__ = _dataclass_getstate
1195-
cls.__setstate__ = _dataclass_setstate
1194+
if '__getstate__' not in cls_dict:
1195+
cls.__getstate__ = _dataclass_getstate
1196+
if '__setstate__' not in cls_dict:
1197+
cls.__setstate__ = _dataclass_setstate
11961198

11971199
return cls
11981200

Lib/test/test_dataclasses.py

+68
Original file line numberDiff line numberDiff line change
@@ -3068,6 +3068,74 @@ def test_frozen_pickle(self):
30683068
self.assertIsNot(obj, p)
30693069
self.assertEqual(obj, p)
30703070

3071+
@dataclass(frozen=True, slots=True)
3072+
class FrozenSlotsGetStateClass:
3073+
foo: str
3074+
bar: int
3075+
3076+
getstate_called: bool = field(default=False, compare=False)
3077+
3078+
def __getstate__(self):
3079+
object.__setattr__(self, 'getstate_called', True)
3080+
return [self.foo, self.bar]
3081+
3082+
@dataclass(frozen=True, slots=True)
3083+
class FrozenSlotsSetStateClass:
3084+
foo: str
3085+
bar: int
3086+
3087+
setstate_called: bool = field(default=False, compare=False)
3088+
3089+
def __setstate__(self, state):
3090+
object.__setattr__(self, 'setstate_called', True)
3091+
object.__setattr__(self, 'foo', state[0])
3092+
object.__setattr__(self, 'bar', state[1])
3093+
3094+
@dataclass(frozen=True, slots=True)
3095+
class FrozenSlotsAllStateClass:
3096+
foo: str
3097+
bar: int
3098+
3099+
getstate_called: bool = field(default=False, compare=False)
3100+
setstate_called: bool = field(default=False, compare=False)
3101+
3102+
def __getstate__(self):
3103+
object.__setattr__(self, 'getstate_called', True)
3104+
return [self.foo, self.bar]
3105+
3106+
def __setstate__(self, state):
3107+
object.__setattr__(self, 'setstate_called', True)
3108+
object.__setattr__(self, 'foo', state[0])
3109+
object.__setattr__(self, 'bar', state[1])
3110+
3111+
def test_frozen_slots_pickle_custom_state(self):
3112+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
3113+
with self.subTest(proto=proto):
3114+
obj = self.FrozenSlotsGetStateClass('a', 1)
3115+
dumped = pickle.dumps(obj, protocol=proto)
3116+
3117+
self.assertTrue(obj.getstate_called)
3118+
self.assertEqual(obj, pickle.loads(dumped))
3119+
3120+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
3121+
with self.subTest(proto=proto):
3122+
obj = self.FrozenSlotsSetStateClass('a', 1)
3123+
obj2 = pickle.loads(pickle.dumps(obj, protocol=proto))
3124+
3125+
self.assertTrue(obj2.setstate_called)
3126+
self.assertEqual(obj, obj2)
3127+
3128+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
3129+
with self.subTest(proto=proto):
3130+
obj = self.FrozenSlotsAllStateClass('a', 1)
3131+
dumped = pickle.dumps(obj, protocol=proto)
3132+
3133+
self.assertTrue(obj.getstate_called)
3134+
3135+
obj2 = pickle.loads(dumped)
3136+
self.assertTrue(obj2.setstate_called)
3137+
self.assertEqual(obj, obj2)
3138+
30713139
def test_slots_with_default_no_init(self):
30723140
# Originally reported in bpo-44649.
30733141
@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)