Merge pull request #50 from msiemens/custom-serialization

Add support for custom serializers

Author: msiemens

docs/extend.rst

@@ -1,6 +1,42 @@
 How to Extend TinyDB
 ====================
 
+Write a Serializer
+------------------
+
+TinyDB's default storage is fairly limited when it comes to supported data types.
+If you need more flexibility, you can implement a Serializer. This allows TinyDB
+to handle classes it otherwise couldn't serialize. Let's see how a Serializer
+for ``datetime`` objects could look like:
+
+.. code-block:: python
+
+    from datetime import datetime
+
+    class DateTimeSerializer(Serializer):
+        OBJ_CLASS = datetime  # The class this serializer handles
+
+        def encode(self, obj):
+            return obj.strftime('%Y-%m-%dT%H:%M:%S')
+
+        def decode(self, s):
+            return datetime.strptime(s, '%Y-%m-%dT%H:%M:%S')
+
+To use the new serializer, we need to use the serialization middleware:
+
+.. code-block:: python
+
+    >>> from tinydb.storages import JSONStorage
+    >>> from tinydb.middlewares import SerializationMiddleware
+    >>>
+    >>> serialization = SerializationMiddleware()
+    >>> serialization.register_serializer(DateTimeSerializer(), 'TinyDate')
+    >>>
+    >>> db = TinyDB('db.json', storage=serialization)
+    >>> db.insert({'date': datetime(2000, 1, 1, 12, 0, 0)})
+    >>> db.all()
+    [{'date': datetime.datetime(2000, 1, 1, 12, 0)}]
+
 Write a custom Storage
 ----------------------
 

docs/usage.rst

@@ -427,8 +427,9 @@ What's next
 Congratulations, you've made through the user guide! Now go and build something
 awesome or dive deeper into TinyDB with these ressources:
 
-- Want to learn how to customize TinyDB and what extensions exist?
-  Check out :doc:`extend`.
+- Want to learn how to customize TinyDB (custom serializers, storages,
+  middlewares) and what extensions exist? Check out :doc:`extend` and
+  :doc:`extensions`.
 - Want to study the API in detail? Read :doc:`api`.
 - Interested in contributing to the TinyDB development guide? Go on to the
   :doc:`contribute`.

tinydb/database.py

@@ -30,6 +30,8 @@ class TinyDB(object):
     and getting tables.
     """
 
+    DEFAULT_STORAGE = JSONStorage
+
     def __init__(self, *args, **kwargs):
         """
         Create a new instance of TinyDB.
@@ -40,7 +42,7 @@ def __init__(self, *args, **kwargs):
         :param storage: The class of the storage to use. Will be initialized
                         with ``args`` and ``kwargs``.
         """
-        storage = kwargs.pop('storage', JSONStorage)
+        storage = kwargs.pop('storage', TinyDB.DEFAULT_STORAGE)
         #: :type: Storage
         self._storage = storage(*args, **kwargs)
 
@@ -153,7 +155,6 @@ def _write_table(self, values, table):
         data[table] = values
 
         self._write(data)
-        return
 
     # Methods that are executed on the default table
     # Because magic methods are not handlet by __getattr__ we need to forward

tinydb/middlewares.py

@@ -2,6 +2,8 @@
 Contains the :class:`base class <tinydb.middlewares.Middleware>` for
 middlewares and implementations.
 """
+from tinydb import TinyDB
+from tinydb.storages import JSONStorage
 
 
 class Middleware(object):
@@ -17,7 +19,7 @@ class Middleware(object):
     example).
     """
 
-    def __init__(self, storage_cls):
+    def __init__(self, storage_cls=TinyDB.DEFAULT_STORAGE):
         self._storage_cls = storage_cls
         self.storage = None
 
@@ -84,7 +86,7 @@ class CachingMiddleware(Middleware):
     #: The number of write operations to cache before writing to disc
     WRITE_CACHE_SIZE = 1000
 
-    def __init__(self, storage_cls):
+    def __init__(self, storage_cls=TinyDB.DEFAULT_STORAGE):
         super(CachingMiddleware, self).__init__(storage_cls)
 
         self.cache = None
@@ -116,3 +118,75 @@ def flush(self):
     def close(self):
         self.flush()
         self.storage.close()
+
+
+class SerializationMiddleware(Middleware):
+    """
+    Provide custom serialization for TinyDB.
+
+    This middleware allows users of TinyDB to register custom serializations.
+    The serialized data will be passed to the wrapped storage and data that
+    is read from the storage will be deserialized.
+    """
+
+    def __init__(self, storage_cls=TinyDB.DEFAULT_STORAGE):
+        super(SerializationMiddleware, self).__init__(storage_cls)
+
+        self._serializers = {}
+
+    def register_serializer(self, serializer, name):
+        """
+        Register a new Serializer.
+
+        When reading from/writing to the underlying storage, TinyDB
+        will run all objects through the list of registered serializers
+        allowing each one to handle objects it recognizes.
+
+        .. note:: The name has to be unique among this database instance.
+                  Re-using the same name will overwrite the old serializer.
+                  Also, registering a serializer will be reflected in all
+                  tables when reading/writing them.
+
+        :param serializer: an instance of the serializer
+        :type serializer: tinydb.serialize.Serializer
+        """
+        self._serializers[name] = serializer
+
+    def read(self):
+        data = self.storage.read()
+
+        for serializer_name in self._serializers:
+            serializer = self._serializers[serializer_name]
+            tag = '{{{}}}:'.format(serializer_name)  # E.g: '{TinyDate}:'
+
+            for eid in data:
+                for field in data[eid]:
+                    try:
+                        if data[eid][field].startswith(tag):
+                            encoded = data[eid][field][len(tag):]
+                            data[eid][field] = serializer.decode(encoded)
+                    except AttributeError:
+                        pass  # Not a string
+
+        return data
+
+    def write(self, data):
+        for serializer_name in self._serializers:
+            # If no serializers are registered, this code will just look up
+            # the serializer list and continue. But if there are serializers,
+            # the inner loop will run very often.
+            # For that reason, the lookup of the serialized class is pulled
+            # out into the outer loop:
+
+            serializer = self._serializers[serializer_name]
+            serializer_class = serializer.OBJ_CLASS
+
+            for eid in data:
+                for field in data[eid]:
+                    if isinstance(data[eid][field], serializer_class):
+                        encoded = serializer.encode(data[eid][field])
+                        tagged = '{{{}}}:{}'.format(serializer_name, encoded)
+
+                        data[eid][field] = tagged
+
+        self.storage.write(data)

tinydb/utils.py

@@ -103,3 +103,19 @@ def catch_warning(warning_cls):
         warnings.filterwarnings('error', category=warning_cls)
 
         yield
+
+
+class AbstractAttribute(object):
+    """An abstract class attribute.
+
+    Use this instead of an abstract property when you don't expect the
+    attribute to be implemented by a property.
+    """
+
+    __isabstractmethod__ = True
+
+    def __init__(self, doc=""):
+        self.__doc__ = doc
+
+    def __get__(self, obj, cls):
+        return self