feihoo87 / waveforms / 7111190278

import functools

import hashlib

import pathlib

import pickle

import time

from typing import Any, Callable, Hashable, KeysView, Optional, Union

try:

    from portalocker import Lock

except:

    import warnings

    warnings.warn('portalocker not installed, cache will not be thread safe!')

    class Lock():

        def __init__(self, file, mode):

            self.file = file

            self.mode = mode

            self.fh = None

        def __enter__(self):

            self.fh = open(self.file, self.mode)

            return self.fh

        def __exit__(self, exc_type, exc_value, traceback):

            self.fh.close()

cache_dir = pathlib.Path.home() / '.waveforms' / 'cache'

cache_dir.mkdir(parents=True, exist_ok=True)

MAXVALUESIZE = 1024

Decorator = Callable[[Callable], Callable]

class Cache(dict):

    def __init__(self, name: str, path: pathlib.Path = cache_dir):

        name = name.split('.')

        self.path: pathlib.Path = path.joinpath(*name[:-1]) / (name[-1] +

        self._buffer = {}

        self._index = {}

        self._mtime = 0

        (self.path / 'values').mkdir(parents=True, exist_ok=True)

    @property

    def index(self) -> pathlib.Path:

        return self.path / 'index'

    @staticmethod

    def _hash(key: Hashable) -> str:

        return hashlib.md5(pickle.dumps(key)).hexdigest()

    def _syncIndex(self) -> None:

        if not self.index.exists() or self._mtime > self.index.stat().st_mtime:

            with Lock(self.index, 'wb') as fh:

                pickle.dump(self._index, fh)

        elif self._mtime < self.index.stat().st_mtime:

            with Lock(self.index, 'rb') as fh:

                self._index = pickle.load(fh)

        self._mtime = self.index.stat().st_mtime

    def _valuePath(self, key: Hashable) -> pathlib.Path:

        hashedKey = self._hash(key)

        return self.path / 'values' / hashedKey

    def _storeValue(self, k: Hashable, buf: bytes) -> None:

        with Lock(self._valuePath(k), 'wb') as fh:

            fh.write(buf)

    def _loadValue(self, k: Hashable) -> None:

        with Lock(self._valuePath(k), 'rb') as fh:

            buf = fh.read()

        hashstr = self._hash(buf)

        self._index[k] = self._valuePath(k).stat().st_mtime, None, hashstr

        self._buffer[k] = pickle.loads(buf)

    def __setitem__(self, k: Hashable, v: Any) -> None:

        buf = pickle.dumps(v)

        if len(buf) <= MAXVALUESIZE:

            mtime = time.time()

            self._index[k] = mtime, buf, ''

            self._mtime = mtime

            self._valuePath(k).unlink(missing_ok=True)

            hashstr = self._hash(buf)

            if k not in self._index or hashstr != self._index[k][2]:

                self._storeValue(k, buf)

                mtime = self._valuePath(k).stat().st_mtime

                self._index[k] = mtime, None, hashstr

                self._mtime = mtime

        self._syncIndex()

        self._buffer[k] = v

    def __getitem__(self, k: Hashable) -> Any:

        self._syncIndex()

        mtime, buf, hashstr = self._index[k]

        if k not in self._buffer and buf is not None:

            self._buffer[k] = pickle.loads(buf)

        elif k not in self._buffer:

            self._loadValue(k)

        return self._buffer[k]

    def __contains__(self, x: Hashable) -> bool:

        self._syncIndex()

        return x in self._index

    def keys(self) -> KeysView[Hashable]:

        self._syncIndex()

        return self._index.keys()

    def clear(self) -> None:

        self._buffer.clear()

        self._index.clear()

        self.index.unlink()

        for f in (self.path / 'values').iterdir():

            f.unlink()

        self._mtime = 0

__caches = {}

def _getCache(name: str) -> Cache:

    try:

        return __caches[name]

    except:

        __caches[name] = Cache(name)

        return __caches[name]

def cache(storage: Optional[Union[str, Cache]] = None) -> Decorator:

    def decorator(func: Callable,

        if storage is None:

            storage = _getCache('.'.join([func.__module__, func.__name__]))

        elif isinstance(storage, str):

            storage = _getCache(storage)

        elif not isinstance(storage, Cache):

            raise Exception(f'storage type {type(storage)} not understand!')

        @functools.wraps(func)

        def wrapper(*args, **kwds):

            kwds = {k: kwds[k] for k in sorted(kwds)}

            try:

                ret = storage[pickle.dumps((args, kwds))]

            except:

                ret = func(*args, **kwds)

                storage[pickle.dumps((args, kwds))] = ret

            return ret

        wrapper.cache = storage

        return wrapper

    return decorator

def clear(name: str) -> None:

    cache = _getCache(name)

    cache.clear()