"""
This module contains the main component of TinyDB: the database.
"""
from typing import Dict, Iterator, Set, Type
from . import JSONStorage
from .storages import Storage
from .table import Table, Document
from .utils import with_typehint
# The table's base class. This is used to add type hinting from the Table
# class to TinyDB. Currently, this supports PyCharm, Pyright/VS Code and MyPy.
TableBase: Type[Table] = with_typehint(Table)
[docs]class TinyDB(TableBase):
"""
The main class of TinyDB.
The ``TinyDB`` class is responsible for creating the storage class instance
that will store this database's documents, managing the database
tables as well as providing access to the default table.
For table management, a simple ``dict`` is used that stores the table class
instances accessible using their table name.
Default table access is provided by forwarding all unknown method calls
and property access operations to the default table by implementing
``__getattr__``.
When creating a new instance, all arguments and keyword arguments (except
for ``storage``) will be passed to the storage class that is provided. If
no storage class is specified, :class:`~tinydb.storages.JSONStorage` will be
used.
.. admonition:: Customization
For customization, the following class variables can be set:
- ``table_class`` defines the class that is used to create tables,
- ``default_table_name`` defines the name of the default table, and
- ``default_storage_class`` will define the class that will be used to
create storage instances if no other storage is passed.
.. versionadded:: 4.0
.. admonition:: Data Storage Model
Data is stored using a storage class that provides persistence for a
``dict`` instance. This ``dict`` contains all tables and their data.
The data is modelled like this::
{
'table1': {
0: {document...},
1: {document...},
},
'table2': {
...
}
}
Each entry in this ``dict`` uses the table name as its key and a
``dict`` of documents as its value. The document ``dict`` contains
document IDs as keys and the documents themselves as values.
:param storage: The class of the storage to use. Will be initialized
with ``args`` and ``kwargs``.
"""
#: The class that will be used to create table instances
#:
#: .. versionadded:: 4.0
table_class = Table
#: The name of the default table
#:
#: .. versionadded:: 4.0
default_table_name = '_default'
#: The class that will be used by default to create storage instances
#:
#: .. versionadded:: 4.0
default_storage_class = JSONStorage
def __init__(self, *args, **kwargs) -> None:
"""
Create a new instance of TinyDB.
"""
storage = kwargs.pop('storage', self.default_storage_class)
# Prepare the storage
self._storage: Storage = storage(*args, **kwargs)
self._opened = True
self._tables: Dict[str, Table] = {}
def __repr__(self):
args = [
'tables={}'.format(list(self.tables())),
'tables_count={}'.format(len(self.tables())),
'default_table_documents_count={}'.format(self.__len__()),
'all_tables_documents_count={}'.format(
['{}={}'.format(table, len(self.table(table)))
for table in self.tables()]),
]
return '<{} {}>'.format(type(self).__name__, ', '.join(args))
[docs] def table(self, name: str, **kwargs) -> Table:
"""
Get access to a specific table.
If the table hasn't been accessed yet, a new table instance will be
created using the :attr:`~tinydb.database.TinyDB.table_class` class.
Otherwise, the previously created table instance will be returned.
All further options besides the name are passed to the table class which
by default is :class:`~tinydb.table.Table`. Check its documentation
for further parameters you can pass.
:param name: The name of the table.
:param kwargs: Keyword arguments to pass to the table class constructor
"""
if name in self._tables:
return self._tables[name]
table = self.table_class(self.storage, name, **kwargs)
self._tables[name] = table
return table
[docs] def tables(self) -> Set[str]:
"""
Get the names of all tables in the database.
:returns: a set of table names
"""
# TinyDB stores data as a dict of tables like this:
#
# {
# '_default': {
# 0: {document...},
# 1: {document...},
# },
# 'table1': {
# ...
# }
# }
#
# To get a set of table names, we thus construct a set of this main
# dict which returns a set of the dict keys which are the table names.
#
# Storage.read() may return ``None`` if the database file is empty,
# so we need to consider this case to and return an empty set in this
# case.
return set(self.storage.read() or {})
[docs] def drop_tables(self) -> None:
"""
Drop all tables from the database. **CANNOT BE REVERSED!**
"""
# We drop all tables from this database by writing an empty dict
# to the storage thereby returning to the initial state with no tables.
self.storage.write({})
# After that we need to remember to empty the ``_tables`` dict, so we'll
# create new table instances when a table is accessed again.
self._tables.clear()
[docs] def drop_table(self, name: str) -> None:
"""
Drop a specific table from the database. **CANNOT BE REVERSED!**
:param name: The name of the table to drop.
"""
# If the table is currently opened, we need to forget the table class
# instance
if name in self._tables:
del self._tables[name]
data = self.storage.read()
# The database is uninitialized, there's nothing to do
if data is None:
return
# The table does not exist, there's nothing to do
if name not in data:
return
# Remove the table from the data dict
del data[name]
# Store the updated data back to the storage
self.storage.write(data)
@property
def storage(self) -> Storage:
"""
Get the storage instance used for this TinyDB instance.
:return: This instance's storage
:rtype: Storage
"""
return self._storage
[docs] def close(self) -> None:
"""
Close the database.
This may be needed if the storage instance used for this database
needs to perform cleanup operations like closing file handles.
To ensure this method is called, the TinyDB instance can be used as a
context manager::
with TinyDB('data.json') as db:
db.insert({'foo': 'bar'})
Upon leaving this context, the ``close`` method will be called.
"""
self._opened = False
self.storage.close()
def __enter__(self):
"""
Use the database as a context manager.
Using the database as a context manager ensures that the
:meth:`~tinydb.database.TinyDB.close` method is called upon leaving
the context.
:return: The current instance
"""
return self
def __exit__(self, *args):
"""
Close the storage instance when leaving a context.
"""
if self._opened:
self.close()
def __getattr__(self, name):
"""
Forward all unknown attribute calls to the default table instance.
"""
return getattr(self.table(self.default_table_name), name)
# Here we forward magic methods to the default table instance. These are
# not handled by __getattr__ so we need to forward them manually here
def __len__(self):
"""
Get the total number of documents in the default table.
>>> db = TinyDB('db.json')
>>> len(db)
0
"""
return len(self.table(self.default_table_name))
def __iter__(self) -> Iterator[Document]:
"""
Return an iterator for the default table's documents.
"""
return iter(self.table(self.default_table_name))
