How to Extend TinyDB¶
There are three main ways to extend TinyDB and modify its behaviour:
- custom storages,
- middlewares, and finally
- 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:
The constructor will receive all arguments passed to TinyDB when creating the database instance (except
storage
which TinyDB itself consumes). In other words callingTinyDB('something', storage=YAMLStorage)
will pass'something'
as an argument toYAMLStorage
.We use
yaml.safe_load
as recommended by the PyYAML documentation when processing data from a potentially untrusted source.If the storage is uninitialized, TinyDB expects the storage to return
None
so it can do any internal initialization that is necessary.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 rundb.close()
on yourTinyDB
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:
- Process every table
- 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:
- You have to use the
super(...)
call as shown in the example. To run your own initialization, add it below thesuper(...)
call. - 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.