Make gitlab objects serializable

lnenov · lnenov · commit 9b59adc8951e · 2017-11-03T07:02:30.000+02:00
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_gitlab.py b/gitlab/tests/test_gitlab.py
@@ -18,6 +18,7 @@
 
 from __future__ import print_function
 
+import pickle
 try:
     import unittest
 except ImportError:
@@ -26,7 +27,6 @@
 from httmock import HTTMock  # noqa
 from httmock import response  # noqa
 from httmock import urlmatch  # noqa
-import pickle
 import six
 
 import gitlab
