making service and session objects serializable

Author: maxcountryman

rauth/service.py

@@ -26,6 +26,8 @@ def process_token_request(r, decoder, *args):
 
 
 class Service(object):
+    __attrs__ = ['name', 'base_url', 'authorize_url']
+
     def __init__(self, name, base_url, authorize_url):
         #: The service name, e.g. 'twitter'.
         self.name = name
@@ -36,6 +38,14 @@ def __init__(self, name, base_url, authorize_url):
         #: The authorization URL.
         self.authorize_url = authorize_url
 
+    def __getstate__(self):
+        return dict((attr, getattr(self, attr, None)) for
+                    attr in self.__attrs__)
+
+    def __setstate__(self, state):
+        for attr, value in state.items():
+            setattr(self, attr, value)
+
 
 class OAuth1Service(Service):
     '''
@@ -116,6 +126,12 @@ class OAuth1Service(Service):
          to :class:`rauth.HmacSha1Signature <HmacSha1Signature>`
     :type signature_obj: :class:`SignatureMethod`
     '''
+    __attrs__ = Service.__attrs__ + ['consumer_key',
+                                     'consumer_secret',
+                                     'request_token_url',
+                                     'access_token_url',
+                                     'session_obj']
+
     def __init__(self,
                  consumer_key,
                  consumer_secret,
@@ -396,6 +412,11 @@ class OAuth2Service(Service):
         :class:`OAuth2Session`
     :type session_obj: :class:`rauth.Session`
     '''
+    __attrs__ = Service.__attrs__ + ['client_id',
+                                     'client_secret',
+                                     'access_token_url',
+                                     'session_obj']
+
     def __init__(self,
                  client_id,
                  client_secret,
@@ -557,6 +578,11 @@ class OflyService(Service):
         `rauth.OflySession`
     :type session_obj: :class:`rauth.Session`
     '''
+    __attrs__ = Service.__attrs__ + ['app_id',
+                                     'app_secret',
+                                     'user_id',
+                                     'session_obj']
+
     def __init__(self,
                  app_id,
                  app_secret,

rauth/session.py

@@ -25,6 +25,8 @@
 
 
 class RauthSession(Session):
+    __attrs__ = Session.__attrs__ + ['service']
+
     def __init__(self, service):
         #: A back reference to a service wrapper, if we're using one.
         self.service = service
@@ -88,6 +90,12 @@ class OAuth1Session(RauthSession):
         `None`.
     :type service: :class:`rauth.Service`
     '''
+    __attrs__ = RauthSession.__attrs__ + ['consumer_key',
+                                          'consumer_secret',
+                                          'access_token',
+                                          'access_token_secret',
+                                          'signature']
+
     VERSION = '1.0'
 
     def __init__(self,
@@ -298,6 +306,10 @@ class OAuth2Session(RauthSession):
         `None`.
     :type service: :class:`rauth.Service`
     '''
+    __attrs__ = RauthSession.__attrs__ + ['client_id',
+                                          'client_secret',
+                                          'access_token']
+
     def __init__(self,
                  client_id,
                  client_secret,
@@ -387,6 +399,10 @@ class OflySession(RauthSession):
         `None`.
     :type service: :class:`rauth.Service`
     '''
+    __attrs__ = RauthSession.__attrs__ + ['app_id',
+                                          'app_secret',
+                                          'user_id']
+
     def __init__(self,
                  app_id,
                  app_secret,

tests/test_service.py

@@ -12,6 +12,7 @@
 from datetime import datetime
 
 import json
+import pickle
 
 
 class MutableDatetime(datetime):
@@ -112,3 +113,10 @@ class RequestMixin(object):
     def test_request(self, method, kwargs):
         r = self.session.request(method, 'foo', **kwargs)
         self.assert_ok(r)
+
+
+class ServiceMixin(object):
+    def test_serialize(self):
+        saved = pickle.loads(pickle.dumps(self.service))
+        for attr in self.service.__attrs__:
+            self.assertEqual(getattr(saved, attr), getattr(self.service, attr))

tests/test_service_oauth1.py

@@ -7,7 +7,7 @@
 '''
 
 from base import RauthTestCase
-from test_service import HttpMixin, RequestMixin
+from test_service import HttpMixin, RequestMixin, ServiceMixin
 
 from rauth.service import OAuth1Service
 from rauth.session import OAUTH1_DEFAULT_TIMEOUT, OAuth1Session
@@ -25,9 +25,11 @@
 import requests
 
 import json
+import pickle
 
 
-class OAuth1ServiceTestCase(RauthTestCase, RequestMixin, HttpMixin):
+class OAuth1ServiceTestCase(RauthTestCase, RequestMixin, ServiceMixin,
+                            HttpMixin):
     consumer_key = '000'
     consumer_secret = '111'
 
@@ -306,3 +308,11 @@ def test_get_auth_session(self):
         self.response.content = resp
         s = self.service.get_auth_session('foo', 'bar')
         self.assertIsInstance(s, OAuth1Session)
+
+    def test_pickle_session(self):
+        session = pickle.loads(pickle.dumps(self.session))
+
+        # Add the fake request back to the session
+        session.request = self.fake_request
+        r = session.request('GET', 'http://example.com/', header_auth=True)
+        self.assert_ok(r)

tests/test_service_oauth2.py

@@ -7,7 +7,7 @@
 '''
 
 from base import RauthTestCase
-from test_service import HttpMixin, RequestMixin
+from test_service import HttpMixin, RequestMixin, ServiceMixin
 
 from rauth.service import OAuth2Service
 from rauth.session import OAUTH2_DEFAULT_TIMEOUT, OAuth2Session
@@ -20,9 +20,11 @@
 import requests
 
 import json
+import pickle
 
 
-class OAuth2ServiceTestCase(RauthTestCase, RequestMixin, HttpMixin):
+class OAuth2ServiceTestCase(RauthTestCase, RequestMixin, ServiceMixin,
+                            HttpMixin):
     client_id = '000'
     client_secret = '111'
     access_token = '123'
@@ -136,3 +138,11 @@ def test_get_auth_session(self):
             'access_token=123&expires_in=3600&refresh_token=456'
         s = self.service.get_auth_session()
         self.assertIsInstance(s, OAuth2Session)
+
+    def test_pickle_session(self):
+        session = pickle.loads(pickle.dumps(self.session))
+
+        # Add the fake request back to the session
+        session.request = self.fake_request
+        r = session.request('GET', 'http://example.com/', bearer_auth=True)
+        self.assert_ok(r)

tests/test_service_ofly.py

@@ -8,7 +8,7 @@
 
 from base import RauthTestCase
 from test_service import (FakeHexdigest, HttpMixin, MutableDatetime,
-                          RequestMixin)
+                          RequestMixin, ServiceMixin)
 
 from rauth.service import OflyService
 from rauth.session import OFLY_DEFAULT_TIMEOUT, OflySession
@@ -22,8 +22,11 @@
 
 import requests
 
+import pickle
 
-class OflyServiceTestCase(RauthTestCase, RequestMixin, HttpMixin):
+
+class OflyServiceTestCase(RauthTestCase, RequestMixin, ServiceMixin,
+                          HttpMixin):
     app_id = '000'
     app_secret = '111'
 
@@ -163,3 +166,14 @@ def test_request_with_bad_hash_meth(self):
     def test_get_auth_session(self):
         s = self.service.get_auth_session('foo')
         self.assertIsInstance(s, OflySession)
+
+    def test_pickle_session(self):
+        session = pickle.loads(pickle.dumps(self.session))
+
+        # Add the fake request back to the session
+        session.request = self.fake_request
+        r = self.session.request('GET',
+                                 'http://example.com/',
+                                 user_id=self.user_id,
+                                 hash_meth='md5')
+        self.assert_ok(r)