From 5def5b35eb1ca9aed6b253f2019a27ce9c863e92 Mon Sep 17 00:00:00 2001 From: Fabio Caccamo Date: Thu, 9 Mar 2023 01:13:24 +0100 Subject: [PATCH] Fix set state when loading from pickle. --- benedict/dicts/base/base_dict.py | 3 ++- benedict/dicts/keyattr/keyattr_dict.py | 4 ++++ benedict/dicts/keypath/keypath_dict.py | 4 ++++ tests/dicts/io/test_pickle.py | 7 ++++++- 4 files changed, 16 insertions(+), 2 deletions(-) diff --git a/benedict/dicts/base/base_dict.py b/benedict/dicts/base/base_dict.py index eb51ebf2..b9de67df 100644 --- a/benedict/dicts/base/base_dict.py +++ b/benedict/dicts/base/base_dict.py @@ -87,7 +87,8 @@ def __setitem__(self, key, value): super().__setitem__(key, value) def __setstate__(self, state): - self.__dict__ = state + self._dict = state["_dict"] + self._pointer = state["_pointer"] def __str__(self): if self._pointer: diff --git a/benedict/dicts/keyattr/keyattr_dict.py b/benedict/dicts/keyattr/keyattr_dict.py index ab8f2765..a1f93512 100644 --- a/benedict/dicts/keyattr/keyattr_dict.py +++ b/benedict/dicts/keyattr/keyattr_dict.py @@ -40,3 +40,7 @@ def __setattr__(self, attr, value): if not self._keyattr_enabled: raise AttributeError self.__setitem__(attr, value) + + def __setstate__(self, state): + super().__setstate__(state) + self._keyattr_enabled = state["_keyattr_enabled"] diff --git a/benedict/dicts/keypath/keypath_dict.py b/benedict/dicts/keypath/keypath_dict.py index 8cd60ff9..d372c144 100644 --- a/benedict/dicts/keypath/keypath_dict.py +++ b/benedict/dicts/keypath/keypath_dict.py @@ -34,6 +34,10 @@ def __setitem__(self, key, value): keypath_util.check_keys(value, self._keypath_separator) super().__setitem__(self._parse_key(key), value) + def __setstate__(self, state): + super().__setstate__(state) + self._keypath_separator = state["_keypath_separator"] + def _parse_key(self, key): keys = keypath_util.parse_keys(key, self._keypath_separator) keys_count = len(keys) diff --git a/tests/dicts/io/test_pickle.py b/tests/dicts/io/test_pickle.py index ce75bb81..63689f04 100644 --- a/tests/dicts/io/test_pickle.py +++ b/tests/dicts/io/test_pickle.py @@ -21,12 +21,17 @@ def test_pickle(self): "h": "0", "i": benedict({"h": True}), } - b = benedict(d, keypath_separator="/") + b = benedict( + d, + keyattr_enabled=False, + keypath_separator="/", + ) b_encoded = pickle.dumps(b) # print(b_encoded) b_decoded = pickle.loads(b_encoded) # print(b_decoded) # print(b_decoded.keypath_separator) self.assertTrue(isinstance(b_decoded, benedict)) + self.assertEqual(b_decoded.keyattr_enabled, b.keyattr_enabled) self.assertEqual(b_decoded.keypath_separator, b.keypath_separator) self.assertEqual(b_decoded, b)