Skip to content

Instantly share code, notes, and snippets.

@initialed85
Last active July 8, 2019 10:53
Show Gist options
  • Save initialed85/79322ee1093d7cc62ae65b0cf099d5b6 to your computer and use it in GitHub Desktop.
Save initialed85/79322ee1093d7cc62ae65b0cf099d5b6 to your computer and use it in GitHub Desktop.
Example of custom Python3 pickling behaviour (should work with Python 2)
from threading import RLock
class Something(object):
def __init__(self, a, b):
self.a = a
self.b = b
self._lock = RLock() # this is not pickleable
class PicklableSomething(Something):
def __getstate__(self): # called when pickled
state = {}
for k, v in self.__dict__.items():
if k == '_lock': # special treatment for _lock property
v = None
state[k] = v
return state
def __setstate__(self, state): # called when unpickled
new_state = {}
for k, v in state.items():
if k == '_lock': # special treatment for _lock property
v = new_state[k] = RLock()
new_state[k] = v
self.__dict__ = new_state
import pickle
import unittest
from something import Something, PicklableSomething
_A = 'Something'
_B = 1337
class SomethingTest(unittest.TestCase):
def setUp(self):
self._subject = Something(_A, _B)
def test_constructor(self):
self.assertEqual(
_A,
self._subject.a,
)
self.assertEqual(
_B,
self._subject.b,
)
def test_pickle(self): # it's expected that pickle fails
with self.assertRaises(TypeError) as cm:
_ = pickle.dumps(self._subject)
self.assertIsInstance(
cm.exception,
TypeError,
)
self.assertEqual(
"can't pickle _thread.RLock objects",
cm.exception.args[0],
)
class PicklableSomethingTest(unittest.TestCase):
def setUp(self):
self._subject = PicklableSomething(_A, _B)
def test_constructor(self):
self.assertEqual(
_A,
self._subject.a,
)
self.assertEqual(
_B,
self._subject.b,
)
def test_pickle(self): # it's expected that pickle succeeds
pickled_subject = pickle.dumps(self._subject)
unpickled_subject = pickle.loads(pickled_subject)
self.assertEqual(
self._subject.a,
unpickled_subject.a,
)
self.assertEqual(
self._subject.b,
unpickled_subject.b,
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment