Merge pull request #99 from pfhayes/serializer

Allow custom serializer objects

beaker/session.py

@@ -1,4 +1,4 @@
-from ._compat import PY2, pickle, http_cookies, unicode_text, b64encode, b64decode
+from ._compat import PY2, pickle, http_cookies, unicode_text, b64encode, b64decode, string_type
 
 import os
 import time
@@ -98,7 +98,9 @@ class Session(dict):
     :param cookie_domain: Domain to use for the cookie.
     :param cookie_path: Path to use for the cookie.
     :param data_serializer: If ``"json"`` or ``"pickle"`` should be used
-                              to serialize data. By default ``pickle`` is used.
+                              to serialize data. Can also be an object with
+                              ``loads` and ``dumps`` methods. By default
+                              ``"pickle"`` is used.
     :param secure: Whether or not the cookie should only be sent over SSL.
     :param httponly: Whether or not the cookie should only be accessible by
                      the browser not by JavaScript.
@@ -140,7 +142,8 @@ def __init__(self, request, id=None, invalidate_corrupt=False,
         self.save_atime = save_accessed_time
         self.use_cookies = use_cookies
         self.cookie_expires = cookie_expires
-        self.data_serializer = data_serializer
+
+        self._set_serializer(data_serializer)
 
         # Default cookie domain/path
         self._domain = cookie_domain
@@ -186,6 +189,17 @@ def __init__(self, request, id=None, invalidate_corrupt=False,
                 else:
                     raise
 
+    def _set_serializer(self, data_serializer):
+        self.data_serializer = data_serializer
+        if self.data_serializer == 'json':
+            self.serializer = util.JsonSerializer()
+        elif self.data_serializer == 'pickle':
+            self.serializer = util.PickleSerializer()
+        elif isinstance(self.data_serializer, string_type):
+            raise BeakerException('Invalid value for data_serializer: %s' % data_serializer)
+        else:
+            self.serializer = data_serializer
+
     def has_key(self, name):
         return name in self
 
@@ -278,10 +292,10 @@ def _encrypt_data(self, session_data=None):
             nonce = b64encode(os.urandom(nonce_len))[:nonce_b64len]
             encrypt_key = crypto.generateCryptoKeys(self.encrypt_key,
                                                     self.validate_key + nonce, 1)
-            data = util.serialize(session_data, self.data_serializer)
+            data = self.serializer.dumps(session_data)
             return nonce + b64encode(crypto.aesEncrypt(data, encrypt_key))
         else:
-            data = util.serialize(session_data, self.data_serializer)
+            data = self.serializer.dumps(session_data)
             return b64encode(data)
 
     def _decrypt_data(self, session_data):
@@ -307,7 +321,7 @@ def _decrypt_data(self, session_data):
             data = b64decode(session_data)
 
         try:
-            return util.deserialize(data, self.data_serializer)
+            return self.serializer.loads(data)
         except:
             if self.invalidate_corrupt:
                 return None
@@ -503,14 +517,15 @@ class CookieSession(Session):
     :param cookie_domain: Domain to use for the cookie.
     :param cookie_path: Path to use for the cookie.
     :param data_serializer: If ``"json"`` or ``"pickle"`` should be used
-                              to serialize data. By default ``pickle`` is used.
+                              to serialize data. Can also be an object with
+                              ``loads` and ``dumps`` methods. By default
+                              ``"pickle"`` is used.
     :param secure: Whether or not the cookie should only be sent over SSL.
     :param httponly: Whether or not the cookie should only be accessible by
                      the browser not by JavaScript.
     :param encrypt_key: The key to use for the local session encryption, if not
                         provided the session will not be encrypted.
     :param validate_key: The key used to sign the local encrypted session
-
     """
     def __init__(self, request, key='beaker.session.id', timeout=None,
                  save_accessed_time=True, cookie_expires=True, cookie_domain=None,
@@ -535,7 +550,8 @@ def __init__(self, request, key='beaker.session.id', timeout=None,
         self.httponly = httponly
         self._domain = cookie_domain
         self._path = cookie_path
-        self.data_serializer = data_serializer
+
+        self._set_serializer(data_serializer)
 
         try:
             cookieheader = request['cookie']

beaker/util.py

@@ -446,15 +446,33 @@ def func_namespace(func):
         return '%s|%s' % (inspect.getsourcefile(func), func.__name__)
 
 
-def serialize(data, method):
-    if method == 'json':
+class PickleSerializer(object):
+    def loads(self, data_string):
+        return pickle.loads(data_string)
+
+    def dumps(self, data):
+        return pickle.dumps(data, 2)
+
+
+class JsonSerializer(object):
+    def loads(self, data_string):
+        return json.loads(zlib.decompress(data_string).decode('utf-8'))
+
+    def dumps(self, data):
         return zlib.compress(json.dumps(data).encode('utf-8'))
+
+
+def serialize(data, serializer):
+    if method == 'json':
+        serializer = JsonSerializer()
     else:
-        return pickle.dumps(data, 2)
+        serializer = PickleSerializer()
+    return serializer.dumps(data)
 
 
 def deserialize(data_string, method):
     if method == 'json':
-        return json.loads(zlib.decompress(data_string).decode('utf-8'))
+        serializer = JsonSerializer()
     else:
-        return pickle.loads(data_string)
+        serializer = PickleSerializer()
+    return serializer.loads(data_string)

tests/test_cookie_only.py

@@ -1,6 +1,7 @@
 import datetime, time
 import re
 import os
+import json
 
 import beaker.session
 import beaker.util
@@ -105,6 +106,36 @@ def test_pickle_serializer():
     res = app.get('/')
     assert 'current value is: 3' in res
 
+def test_custom_serializer():
+    was_used = [False, False]
+    class CustomSerializer(object):
+        def loads(self, data_string):
+            was_used[0] = True
+            return json.loads(data_string.decode('utf-8'))
+
+        def dumps(self, data):
+            was_used[1] = True
+            return json.dumps(data).encode('utf-8')
+
+    serializer = CustomSerializer()
+    options = {'session.validate_key':'hoobermas', 'session.type':'cookie', 'data_serializer': serializer}
+    app = TestApp(SessionMiddleware(simple_app, **options))
+
+    res = app.get('/')
+    assert 'current value is: 1' in res
+
+    res = app.get('/')
+    cookie = SignedCookie('hoobermas')
+    session_data = cookie.value_decode(app.cookies['beaker.session.id'])[0]
+    session_data = b64decode(session_data)
+    data = serializer.loads(session_data)
+    assert data['value'] == 2
+
+    res = app.get('/')
+    assert 'current value is: 3' in res
+
+    assert all(was_used)
+
 def test_expires():
     options = {'session.validate_key':'hoobermas', 'session.type':'cookie',
                'session.cookie_expires': datetime.timedelta(days=1)}