How to Extend TinyDB

There are three main ways to extend TinyDB and modify its behaviour:

  1. custom storages,
  2. middlewares, and finally
  3. custom table classes.

Let’s look at them in this order.

Write a Custom Storage

First, we have support for custom storages. By default TinyDB comes with an in-memory storage and a JSON file storage. But of course you can add your own. Let’s look how you could add a YAML storage using PyYAML:

import yaml

class YAMLStorage(Storage):
    def __init__(self, filename):  # (1)
        self.filename = filename

    def read(self):
        with open(self.filename) as handle:
            try:
                data = yaml.safe_load(handle.read())  # (2)
                return data
            except yaml.YAMLError:
                return None  # (3)

    def write(self, data):
        with open(self.filename, 'w') as handle:
            yaml.dump(data, handle)

    def close(self):  # (4)
        pass

There are some things we should look closer at:

  1. The constructor will receive all arguments passed to TinyDB when creating the database instance (except storage which TinyDB itself consumes). In other words calling TinyDB('something', storage=YAMLStorage) will pass 'something' as an argument to YAMLStorage.

  2. We use yaml.safe_load as recommended by the PyYAML documentation when processing data from a potentially untrusted source.

  3. If the storage is uninitialized, TinyDB expects the storage to return None so it can do any internal initialization that is necessary.

  4. If your storage needs any cleanup (like closing file handles) before an instance is destroyed, you can put it in the close() method. To run these, you’ll either have to run db.close() on your TinyDB instance or use it as a context manager, like this:

    with TinyDB('db.yml', storage=YAMLStorage) as db:
        # ...
    

Finally, using the YAML storage is very straight-forward:

db = TinyDB('db.yml', storage=YAMLStorage)
# ...

Write a Custom Middleware

Sometimes you don’t want to write a new storage but rather modify the behaviour of an existing one. As an example we’ll build a middleware that filters out any empty items.

Because middlewares act as a wrapper around a storage, they needs a read() and a write(data) method. In addition, they can access the underlying storage via self.storage. Before we start implementing we should look at the structure of the data that the middleware receives. Here’s what the data that goes through the middleware looks like:

{
    '_default': {
        1: {'key': 'value'},
        2: {'key': 'value'},
        # other items
    },
    # other tables
}

Thus, we’ll need two nested loops:

  1. Process every table
  2. Process every item

Now let’s implement that:

class RemoveEmptyItemsMiddleware(Middleware):
    def __init__(self, storage_cls=TinyDB.DEFAULT_STORAGE):
        # Any middleware *has* to call the super constructor
        # with storage_cls
        super(CustomMiddleware, self).__init__(storage_cls)

    def read(self):
        data = self.storage.read()

        for table_name in data:
            table = data[table_name]

            for element_id in table:
                item = table[element_id]

                if item == {}:
                    del table[element_id]

        return data

    def write(self, data):
        for table_name in data:
            table = data[table_name]

            for element_id in table:
                item = table[element_id]

                if item == {}:
                    del table[element_id]

        self.storage.write(data)

    def close(self):
        self.storage.close()

Two remarks:

  1. You have to use the super(...) call as shown in the example. To run your own initialization, add it below the super(...) call.
  2. This is an example for a middleware, not an example for clean code. Don’t use it as shown here without at least refactoring the loops into a separate method.

To wrap a storage with this new middleware, we use it like this:

db = TinyDB(storage=RemoveEmptyItemsMiddleware(SomeStorageClass))

Here SomeStorageClass should be replaced with the storage you want to use. If you leave it empty, the default storage will be used (which is the JSONStorage).

Creating a Custom Table Classes

Custom storages and middlewares are useful if you want to modify the way TinyDB stores its data. But there are cases where you want to modify how TinyDB itself behaves. For that use case TinyDB supports custom table classes. Internally TinyDB creates a Table instance for every table that is used. You can overwrite which class is used by setting TinyDB.table_class before creating a TinyDB instance. This class has to support the Table API. The best way to accomplish that is to subclass it:

from tinydb.database import Table

class YourTableClass(Table):
    pass  # Modify original methods as needed

For an more advanced example, see the source of the tinydb-smartcache extension.