extensions for python tinydb #tinydb

tinydb_atomic_storage.py

# -*- coding: utf-8 -*-
#
# Copyright (c) 2020~2999 - Cologler <[email protected]>
# ----------
# tinydb storage with atomic write.
# ----------

import os
import json
from contextlib import suppress
from abc import abstractmethod

from tinydb.storages import Storage
from atomicwrites import atomic_write

class AtomicStorage(Storage):
    def __init__(self, path: str, create_dirs=False, **kwargs):
        super().__init__()
        self._path = path
        self._create_dirs = create_dirs
        self._kwargs = kwargs

    def read(self):
        buf = None
        if os.path.isfile(self._path):
            with suppress(FileNotFoundError):
                with open(self._path, mode='rb', encoding=encoding) as fp:
                    buf = fp.read()
        return self._load(buf) if buf is not None else None

    def write(self, data: dict):
        if self._create_dirs:
            base_dir = os.path.dirname(path)
            if not os.path.isdir(base_dir):
                os.makedirs(base_dir)

        serialized = self._dump(data)
        with atomic_write(self._path, mode='wb', overwrite=True) as fp:
            fp.write(serialized)

    @abstractmethod
    def _load(self, data: bytes):
        raise NotImplementedError

    @abstractmethod
    def _dump(self, obj) -> bytes:
        raise NotImplementedError


class JSONAtomicStorage(AtomicStorage):

    def __init__(self, encoding=None, **kwargs):
        super().__init__(**kwargs)
        self._encoding = encoding or 'utf-8'

    def _load(self, data: bytes):
        return json.loads(data, encoding=self._encoding)

    def _dump(self, obj) -> bytes:
        return json.dumps(obj, **self._kwargs).encode(self._encoding)

tinydb_filelock_middleware.py

# -*- coding: utf-8 -*-
#
# Copyright (c) 2020~2999 - Cologler <[email protected]>
# ----------
# tinydb middleware with filelock.
# ----------

from tinydb.middlewares import Middleware
from filelock import FileLock

class FileLockMiddleware(Middleware):
    def __init__(self, storage_cls):
        super().__init__(storage_cls)
        self._lock = None

    def __call__(self, *args, **kwargs):
        lockfile = kwargs.pop('lockfile', None)
        if lockfile is None:
            path = kwargs.get('path', None)
            if path is not None:
                lockfile = path + '.lock'
        if not lockfile:
            raise ValueError('missing path or lockfile argument.')
        self._lock = FileLock(lockfile)
        self._lock.acquire()
        return super().__call__(*args, **kwargs)

    def close(self):
        self.storage.close()
        self._lock.release()