You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
274 lines
8.5 KiB
274 lines
8.5 KiB
""" |
|
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))
|
|
|