extend.rst

How to Extend TinyDB

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

1. custom storages,
2. custom middlewares,
3. use hooks and overrides, and
4. subclassing TinyDB and Table.

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 Custom Middleware

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

Because middleware acts 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):
        # Any middleware *has* to call the super constructor
        # with storage_cls
        super().__init__(storage_cls)  # (1)

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

        for table_name in data:
            table_data = data[table_name]

            for doc_id in table_data:
                item = table_data[doc_id]

                if item == {}:
                    del table_data[doc_id]

        return data

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

            for doc_id in table:
                item = table_data[doc_id]

                if item == {}:
                    del table_data[doc_id]

        self.storage.write(data)

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

Note that the constructor calls the middleware constructor (1) and passes the storage class to the middleware constructor.

To wrap 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).

Use hooks and overrides

There are cases when neither creating a custom storage nor using a custom middleware will allow you to adapt TinyDB in the way you need. In this case you can modify TinyDB's behavior by using predefined hooks and override points. For example you can configure the name of the default table by setting TinyDB.default_table_name:

TinyDB.default_table_name = 'my_table_name'

Both :class:`~tinydb.database.TinyDB` and the :class:`~tinydb.table.Table` classes allow modifying their behavior using hooks and overrides. To use Table's overrides, you can access the class using TinyDB.table_class:

TinyDB.table_class.default_query_cache_capacity = 100

Read the :ref:`api_docs` for more details on the available hooks and override points.

Subclassing TinyDB and Table

Finally, there's the last option to modify TinyDB's behavior. That way you can change how TinyDB itself works more deeply than using the other extension mechanisms.

When creating a subclass you can use it by using hooks and overrides to override the default classes that TinyDB uses:

class MyTable(Table):
    # Add your method overrides
    ...

TinyDB.table_class = MyTable

# Continue using TinyDB as usual

TinyDB's source code is documented with extensions in mind, explaining how everything works even for internal methods and classes. Feel free to dig into the source and adapt everything you need for your projects.