"""Contains the :class:`base class <tinydb.storages.Storage>` for storages andimplementations."""importioimportjsonimportosimportwarningsfromabcimportABC,abstractmethodfromtypingimportDict,Any,Optional__all__=('Storage','JSONStorage','MemoryStorage')deftouch(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.    """ifcreate_dirs:base_dir=os.path.dirname(path)# Check if we need to create missing parent directoriesifnotos.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 contentswithopen(path,'a'):pass
[docs]classStorage(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]@abstractmethoddefread(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.        """raiseNotImplementedError('To be overridden!')
[docs]@abstractmethoddefwrite(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.        """raiseNotImplementedError('To be overridden!')
[docs]defclose(self)->None:"""        Optional: Close open file handles, etc.        """pass
[docs]classJSONStorage(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_modeself.kwargs=kwargsifaccess_modenotin('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 modeifany([characterinself._modeforcharacterin('+','w','a')]):# any of the writing modestouch(path,create_dirs=create_dirs)# Open the file for reading/writingself._handle=open(path,mode=self._mode,encoding=encoding)
[docs]defclose(self)->None:self._handle.close()
[docs]defread(self)->Optional[Dict[str,Dict[str,Any]]]:# Get the file size by moving the cursor to the file end and reading# its locationself._handle.seek(0,os.SEEK_END)size=self._handle.tell()ifnotsize:# File is empty, so we return ``None`` so TinyDB can properly# initialize the databasereturnNoneelse:# Return the cursor to the beginning of the fileself._handle.seek(0)# Load the JSON contents of the filereturnjson.load(self._handle)
[docs]defwrite(self,data:Dict[str,Dict[str,Any]]):# Move the cursor to the beginning of the file just in caseself._handle.seek(0)# Serialize the database state using the user-provided argumentsserialized=json.dumps(data,**self.kwargs)# Write the serialized data to the filetry:self._handle.write(serialized)exceptio.UnsupportedOperation:raiseIOError('Cannot write to the database. Access mode is "{0}"'.format(self._mode))# Ensure the file has been writtenself._handle.flush()os.fsync(self._handle.fileno())# Remove data that is behind the new cursor in case the file has# gotten shorterself._handle.truncate()
[docs]classMemoryStorage(Storage):"""    Store the data as JSON in memory.    """
[docs]def__init__(self):"""        Create a new instance.        """super().__init__()self.memory=None
[docs]defread(self)->Optional[Dict[str,Dict[str,Any]]]:returnself.memory
[docs]defwrite(self,data:Dict[str,Dict[str,Any]]):self.memory=data