275 lines
8.5 KiB
Python
275 lines
8.5 KiB
Python
|
"""
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
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))
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
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 {})
|
||
|
|
|
||
|
|
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()
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
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))
|