Source code for tinydb.storages
"""
Contains the :class:`base class <tinydb.storages.Storage>` for storages and
implementations.
"""
import io
import json
import os
import warnings
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional
__all__ = ('Storage', 'JSONStorage', 'MemoryStorage')
def touch(path: str, create_dirs: bool):
"""
Create a file if it doesn't exist yet.
:param path: The file to create.
:param create_dirs: Whether to create all missing parent directories.
"""
if create_dirs:
base_dir = os.path.dirname(path)
# Check if we need to create missing parent directories
if not os.path.exists(base_dir):
os.makedirs(base_dir)
# Create the file by opening it in 'a' mode which creates the file if it
# does not exist yet but does not modify its contents
with open(path, 'a'):
pass
[docs]class Storage(ABC):
"""
The abstract base class for all Storages.
A Storage (de)serializes the current state of the database and stores it in
some place (memory, file on disk, ...).
"""
# Using ABCMeta as metaclass allows instantiating only storages that have
# implemented read and write
[docs] @abstractmethod
def read(self) -> Optional[Dict[str, Dict[str, Any]]]:
"""
Read the current state.
Any kind of deserialization should go here.
Return ``None`` here to indicate that the storage is empty.
"""
raise NotImplementedError('To be overridden!')
[docs] @abstractmethod
def write(self, data: Dict[str, Dict[str, Any]]) -> None:
"""
Write the current state of the database to the storage.
Any kind of serialization should go here.
:param data: The current state of the database.
"""
raise NotImplementedError('To be overridden!')
[docs] def close(self) -> None:
"""
Optional: Close open file handles, etc.
"""
pass
[docs]class JSONStorage(Storage):
"""
Store the data in a JSON file.
"""
[docs] def __init__(self, path: str, create_dirs=False, encoding=None, access_mode='r+', **kwargs):
"""
Create a new instance.
Also creates the storage file, if it doesn't exist and the access mode
is appropriate for writing.
Note: Using an access mode other than `r` or `r+` will probably lead to
data loss or data corruption!
:param path: Where to store the JSON data.
:param access_mode: mode in which the file is opened (r, r+)
:type access_mode: str
"""
super().__init__()
self._mode = access_mode
self.kwargs = kwargs
if access_mode not in ('r', 'rb', 'r+', 'rb+'):
warnings.warn(
'Using an `access_mode` other than \'r\', \'rb\', \'r+\' '
'or \'rb+\' can cause data loss or corruption'
)
# Create the file if it doesn't exist and creating is allowed by the
# access mode
if any([character in self._mode for character in ('+', 'w', 'a')]): # any of the writing modes
touch(path, create_dirs=create_dirs)
# Open the file for reading/writing
self._handle = open(path, mode=self._mode, encoding=encoding)
[docs] def close(self) -> None:
self._handle.close()
[docs] def read(self) -> Optional[Dict[str, Dict[str, Any]]]:
# Get the file size by moving the cursor to the file end and reading
# its location
self._handle.seek(0, os.SEEK_END)
size = self._handle.tell()
if not size:
# File is empty, so we return ``None`` so TinyDB can properly
# initialize the database
return None
else:
# Return the cursor to the beginning of the file
self._handle.seek(0)
# Load the JSON contents of the file
return json.load(self._handle)
[docs] def write(self, data: Dict[str, Dict[str, Any]]):
# Move the cursor to the beginning of the file just in case
self._handle.seek(0)
# Serialize the database state using the user-provided arguments
serialized = json.dumps(data, **self.kwargs)
# Write the serialized data to the file
try:
self._handle.write(serialized)
except io.UnsupportedOperation:
raise IOError('Cannot write to the database. Access mode is "{0}"'.format(self._mode))
# Ensure the file has been written
self._handle.flush()
os.fsync(self._handle.fileno())
# Remove data that is behind the new cursor in case the file has
# gotten shorter
self._handle.truncate()
[docs]class MemoryStorage(Storage):
"""
Store the data as JSON in memory.
"""
[docs] def __init__(self):
"""
Create a new instance.
"""
super().__init__()
self.memory = None
[docs] def read(self) -> Optional[Dict[str, Dict[str, Any]]]:
return self.memory
[docs] def write(self, data: Dict[str, Dict[str, Any]]):
self.memory = data