diff --git a/CHANGELOG.md b/CHANGELOG.md index d6a69db..2f17896 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,12 @@ Pillar Python SDK changelog =========================== +Version 1.7 (in development) +---------------------------- + +- Added support pickling/unpickling resources. + + Version 1.6 ----------- diff --git a/pillarsdk/resource.py b/pillarsdk/resource.py index f1b10e0..83d1ad2 100644 --- a/pillarsdk/resource.py +++ b/pillarsdk/resource.py @@ -140,6 +140,53 @@ class Resource(object): return cls(dict_or_resource) + def __getstate__(self): + """Returns a state suitable for pickling. + + This is basically a copy of to_dict(), except that subobjects + are pickled instead of stored as sub-dict. + """ + import pickle + + def parse_object(value): + if isinstance(value, Resource): + return '__pickled__', pickle.dumps(value) + elif isinstance(value, list): + new_list = [] + for obj in value: + new_list.append(parse_object(obj)) + return new_list + else: + return value + + data = {} + for key in self.__data__: + data[key] = parse_object(self.__data__[key]) + return data + + def __setstate__(self, state): + import pickle + + def is_pickled(subval): + return (isinstance(subval, tuple) + and len(subval) == 2 + and subval[0] == '__pickled__' + and isinstance(subval[1], bytes)) + + for key, val in state.items(): + if is_pickled(val): + state[key] = pickle.loads(val[1]) + elif isinstance(val, list): + for idx, subval in enumerate(val): + if is_pickled(subval): + val[idx] = pickle.loads(subval[1]) + elif isinstance(val, dict): + for subkey, subval in val.items(): + if is_pickled(subval): + val[subkey] = pickle.loads(subval[1]) + + self.__init__(state) + class Find(Resource): @classmethod diff --git a/tests/test_pickle.py b/tests/test_pickle.py new file mode 100644 index 0000000..7d7837a --- /dev/null +++ b/tests/test_pickle.py @@ -0,0 +1,79 @@ +# -*- coding: utf8 -*- +import pickle +import unittest + + +class PickleTest(unittest.TestCase): + def test_pickling_node(self): + from pillarsdk import Node, File + + base_link = 'https://storage.googleapis.com/57534c07c37a7195f/_%2F326760eb7d7b244afe52' + picture = File({ + '_id': '57534ccdc379cf1b24a7196d', + '_created': '2016-06-04T23:49:01.000+0200', + '_updated': '2017-09-07T15:55:37.000+0200', + '_etag': '4bb3f525ea637c612fca882e5fcf334c056544fd', + 'status': 'complete', + 'name': '326763a705364fe0b0eb7d7b244afe52.png', + 'backend': 'gcs', + 'format': 'png', + 'variations': [ + {'width': 160, + 'length': 4705, + 'content_type': 'image/jpeg', + 'height': 160, + 'file_path': '326763a705364fe0b0eb7d7b244afe52-b.jpg', + 'size': 'b', + 'link': base_link + '-b.jpg'}, + {'width': 269, + 'length': 8508, + 'content_type': 'image/jpeg', + 'height': 269, + 'file_path': '326763a705364fe0b0eb7d7b244afe52-h.jpg', + 'size': 'h', + 'link': base_link + '-h.jpg'}, + {'width': 269, + 'length': 8508, + 'content_type': 'image/jpeg', + 'height': 269, + 'file_path': '326763a705364fe0b0eb7d7b244afe52-m.jpg', + 'size': 'm', + 'link': base_link + '-m.jpg'}, + ], + 'filename': '01d.png', + 'project': '57534c07c379cf1b24a7195f', + 'width': 269, + 'length': 9681, + 'user': '573dff22c379cf12e649f07a', + 'content_type': 'image/png', + 'height': 269, + 'file_path': '326763a705364fe0b0eb7d7b244afe52.png', + 'md5': '', + 'length_aggregate_in_bytes': 47050, + 'link': base_link + '.png', + 'link_expires': '2117-09-08T14:54:35.250+0200', + } + ) + parent_node = Node({ + '_id': '54134', + 'name': 'Dadday', + }) + original = Node({ + '_id': '123456', + 'name': 'über cooole node', + 'parent': parent_node, + 'picture': picture, + }) + + pickled = pickle.dumps(original) + restored = pickle.loads(pickled) + + self.assertEqual(restored._id, '123456') + self.assertEqual(restored.name, 'über cooole node') + + self.assertIsInstance(restored.parent, Node) + self.assertEqual(restored.parent.name, 'Dadday') + + url = 'https://storage.googleapis.com/57534c07c37a7195f/_%2F326760eb7d7b244afe52-m.jpg' + self.assertEqual(url, original.picture.thumbnail('m')) + self.assertEqual(url, restored.picture.thumbnail('m'))