Module's base objects serialization (#359)

lnenov · Gauvain Pocentek · commit 226e6ce9e521 · 2017-11-03T15:05:17.000+01:00
Make gitlab objects serializable

With current implementation of API v3 and v4 support, some instances
have properties of type module and are not serializable. Handle
these properties manually with setstate and getstate methods.
diff --git a/gitlab/__init__.py b/gitlab/__init__.py
@@ -138,6 +138,17 @@ def __init__(self, url, private_token=None, email=None, password=None,
                     manager = getattr(objects, cls_name)(self)
                     setattr(self, var_name, manager)
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop('_objects')
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        objects = importlib.import_module('gitlab.v%s.objects' %
+                                          self._api_version)
+        self._objects = objects
+
     @property
     def api_version(self):
         return self._api_version
diff --git a/gitlab/base.py b/gitlab/base.py
@@ -416,6 +416,17 @@ def __init__(self, gl, data=None, **kwargs):
         if not hasattr(self, "id"):
             self.id = None
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        module = state.pop('_module')
+        state['_module_name'] = module.__name__
+        return state
+
+    def __setstate__(self, state):
+        module_name = state.pop('_module_name')
+        self.__dict__.update(state)
+        self._module = importlib.import_module(module_name)
+
     def _set_manager(self, var, cls, attrs):
         manager = cls(self.gitlab, self, attrs)
         setattr(self, var, manager)
@@ -555,6 +566,17 @@ def __init__(self, manager, attrs):
         self.__dict__['_parent_attrs'] = self.manager.parent_attrs
         self._create_managers()
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        module = state.pop('_module')
+        state['_module_name'] = module.__name__
+        return state
+
+    def __setstate__(self, state):
+        module_name = state.pop('_module_name')
+        self.__dict__.update(state)
+        self._module = importlib.import_module(module_name)
+
     def __getattr__(self, name):
         try:
             return self.__dict__['_updated_attrs'][name]
diff --git a/gitlab/tests/test_base.py b/gitlab/tests/test_base.py
@@ -15,6 +15,7 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import pickle
 try:
     import unittest
 except ImportError:
@@ -86,6 +87,15 @@ def test_instanciate(self):
         self.assertEqual(self.manager, obj.manager)
         self.assertEqual(self.gitlab, obj.manager.gitlab)
 
+    def test_pickability(self):
+        obj = FakeObject(self.manager, {'foo': 'bar'})
+        original_obj_module = obj._module
+        pickled = pickle.dumps(obj)
+        unpickled = pickle.loads(pickled)
+        self.assertIsInstance(unpickled, FakeObject)
+        self.assertTrue(hasattr(unpickled, '_module'))
+        self.assertEqual(unpickled._module, original_obj_module)
+
     def test_attrs(self):
         obj = FakeObject(self.manager, {'foo': 'bar'})
 
diff --git a/gitlab/tests/test_gitlab.py b/gitlab/tests/test_gitlab.py
@@ -18,6 +18,7 @@
 
 from __future__ import print_function
 
+import pickle
 try:
     import unittest
 except ImportError:
@@ -890,6 +891,14 @@ def setUp(self):
                          email="testuser@test.com", password="testpassword",
                          ssl_verify=True)
 
+    def test_pickability(self):
+        original_gl_objects = self.gl._objects
+        pickled = pickle.dumps(self.gl)
+        unpickled = pickle.loads(pickled)
+        self.assertIsInstance(unpickled, Gitlab)
+        self.assertTrue(hasattr(unpickled, '_objects'))
+        self.assertEqual(unpickled._objects, original_gl_objects)
+
     def test_credentials_auth_nopassword(self):
         self.gl.email = None
         self.gl.password = None
diff --git a/gitlab/tests/test_gitlabobject.py b/gitlab/tests/test_gitlabobject.py
@@ -21,6 +21,7 @@
 from __future__ import absolute_import
 
 import json
+import pickle
 try:
     import unittest
 except ImportError:
@@ -158,6 +159,15 @@ def test_json(self):
         self.assertEqual(data["username"], "testname")
         self.assertEqual(data["gitlab"]["url"], "http://localhost/api/v3")
 
+    def test_pickability(self):
+        gl_object = CurrentUser(self.gl, data={"username": "testname"})
+        original_obj_module = gl_object._module
+        pickled = pickle.dumps(gl_object)
+        unpickled = pickle.loads(pickled)
+        self.assertIsInstance(unpickled, CurrentUser)
+        self.assertTrue(hasattr(unpickled, '_module'))
+        self.assertEqual(unpickled._module, original_obj_module)
+
     def test_data_for_gitlab(self):
         class FakeObj1(GitlabObject):
             _url = '/fake1'
