Module's base objects serialization (#359)

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.

5 files changed, 62 insertions, 0 deletions

diff --git a/gitlab/__init__.py b/gitlab/__init__.py
index b721030..965bf38 100644
--- a/gitlab/__init__.py
+++ b/gitlab/__init__.py
@@ -138,6 +138,17 @@ class Gitlab(object):
                     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
index 4213d2f..ec5f698 100644
--- a/gitlab/base.py
+++ b/gitlab/base.py
@@ -416,6 +416,17 @@ class GitlabObject(object):
         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 @@ class RESTObject(object):
         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
index c55f000..31dd967 100644
--- 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 @@ class TestRESTObject(unittest.TestCase):
         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
index 0f39624..027de0c 100644
--- 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 @@ class TestGitlab(unittest.TestCase):
                          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
index 695f900..f7fd187 100644
--- a/gitlab/tests/test_gitlabobject.py
+++ b/gitlab/tests/test_gitlabobject.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 from __future__ import absolute_import
 
 import json
+import pickle
 try:
     import unittest
 except ImportError:
@@ -158,6 +159,15 @@ class TestGitlabObject(unittest.TestCase):
         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'