werkzeug cache additions

cache_add.py

'''This module contains cache-related logic'''
## TODO:  improve documentation of contents

import threading
import exceptions
from werkzeug.contrib.cache import (SimpleCache, MemcachedCache, RedisCache,
                                    FileSystemCache, NullCache)


class CacheMixin:
    '''A mix-in for werkzeug.contrib.cache classes to make
    some their methods thread-safe using threading.Lock
    '''
    def __init__(self, *args, **kwargs):
        self._lock = threading.Lock()
        super().__init__(*args, **kwargs)

    def add(self, *args, **kwargs):
        with self._lock:
            return super().add(*args, **kwargs)

    def get(self, *args, **kwargs):
        with self._lock:
            return super().get(*args, **kwargs)

    def set(self, *args, **kwargs):
        with self._lock:
            return super().set(*args, **kwargs)

    def delete(self, *args, **kwargs):
        with self._lock:
            return super().delete(*args, **kwargs)

    def get_strict(self, key):
        with self._lock:
            ret = super().get(key)
            if ret is None and not super().has(key):
                raise exceptions.CacheNotFound
            return ret

    def get_nonstrict(self, key):
        with self._lock:
            ret = super().get(key)
            if ret is None:
                raise exceptions.CacheNotFound
            return ret


# Overwrites for werkzeug.cache.contrib classes to make them thread-safe
class SimpleCache(CacheMixin, SimpleCache):
    pass

class MemcachedCache(CacheMixin, MemcachedCache):
    pass

class RedisCache(CacheMixin, RedisCache):
    pass

class FileSystemCache(CacheMixin, FileSystemCache):
    pass


class HashedArgs:
    __slots__ = '_hash', '_seq'

    def __init__(self, func, args=None, kwargs=None):
        seq = [func]
        if args is not None:
            seq += args
        if kwargs is not None:
            # flattern dict to seq(k1, v1, k2, v2, ...)
            seq += [it for k in sorted(kwargs) for it in [k, kwargs[k]]]
        seq = tuple(seq)
        self._seq = seq
        self._hash = hash(seq)

    def __hash__(self):
        return self._hash

    def __eq__(self, other):
        return isinstance(other, self.__class__) and self._seq == other._seq


class cachedproperty(property):
    '''A property with caching
    Depends on underlying cache, which should be subclass or have interface of
    werkzeug.contrib.cache.BaseCache. Underlying cache should be provided at
    creation time (using `cache` kwarg) or defined at client class as its
    `__cache__` attribute.
    '''

    __slots__ = 'fget', 'fset', 'fdel', 'timeout', 'cache', 'isstrict', 'doc'

    def __init__(self, fget=None, fset=None, fdel=None, doc=None,
                 *, timeout=None, cache=None, isstrict=True):
        self.fget = fget
        self.fset = fset
        self.fdel = fdel
        self.timeout = timeout
        self.cache = cache
        self.isstrict = isstrict
        self.doc = doc if doc is not None else fget.__doc__

    def __get__(self, instance, owner=None):
        if instance is None:  # access as class attribute
            return self
        if self.fget is None:
            raise AttributeError("can't get attribute")
        cache = self.cache or instance.__cache__  # short evaluation here
        timeout = self.timeout or cache.default_timeout
        get = cache.get_strict if self.isstrict else cache.get_nonstrict
        try:
            res = get(self.fget)  # try to get
        except exceptions.CacheNotFound:  # if invalid -- fget and put to cache
            res = self.fget(instance)
            r = cache.set(self.fget, res, timeout=timeout)
            if not r:
                raise exceptions.CacheError from None
        return res

    def __set__(self, instance, value):
        if self.fset is None:
            raise AttributeError("can't set attribute")
        cache = self.cache or instance.__cache__
        cache.delete(self.fget)  # invalidate and run fset func
        self.fset(instance, value)

    def __delete__(self, instance):
        if self.fdel is None:
            raise AttributeError("can't delete attribute")
        cache = self.cache or instance.__cache__
        cache.delete(self.fget)  # invalidate and run fdel func
        self.fdel(instance)


class cachedmethod:
    __slots__ = 'method', 'cache', 'timeout', 'isstrict'

    class cachedmethod_wrapper:
        __slots__ = 'outer', 'instance'
        def __init__(self, outer, instance):
            self.outer = outer
            self.instance = instance

        def __call__(self, *args, **kwargs):
            '''Returns cached ret value of wrapped method'''
            cache = self.outer.cache or self.instance.__cache__
            timeout = self.outer.timeout or cache.default_timeout
            args = (self.instance,) + args
            call_hash = HashedArgs(self.outer.method, args, kwargs)
            get = (cache.get_strict if self.outer.isstrict
                                    else cache.get_nonstrict)
            try:
                ret = get(call_hash)
            except exceptions.CacheNotFound:
                ret = self.outer.method(*args, **kwargs)
                status = cache.set(call_hash, ret, timeout=timeout)
                if not status:
                    raise exceptions.CacheError from None
            return ret

        def cache_push(self, *args, **kwargs):
            '''Perform recalc & push ret to cache for provided args'''
            cache = self.outer.cache or self.instance.__cache__
            timeout = self.outer.timeout or cache.default_timeout
            args = (self.instance,) + args
            call_hash = HashedArgs(self.outer.method, args, kwargs)
            ret = self.outer.method(*args, **kwargs)
            status = cache.set(call_hash, ret, timeout=timeout)
            if not status:
                raise exceptions.CacheError from None
            return ret

        def cache_bypass(self, *args, **kwargs):
            '''Bypasses cachig mechanism and acts like normal method call'''
            args = (self.instance,) + args
            return self.outer.method(*args, **kwargs)

        def cache_invalidate(self, *args, **kwargs):
            cache = self.outer.cache or self.instance.__cache__
            args = (self.instance,) + args
            call_hash = HashedArgs(self.outer.method, args, kwargs)
            return cache.delete(call_hash)

    # descriptor methods below
    def __init__(self, timeout=None, cache=None, isstrict=True):
        # called to create decorator, like: @cachedmethod(cache=...)
        self.cache = cache
        self.timeout = timeout
        self.isstrict = isstrict

    def __get__(self, instance, owner=None):
        if instance is None:  # if get as class attr, not through instance
            return self
        return self.cachedmethod_wrapper(outer=self, instance=instance)

    def __call__(self, method):
        # a call to decorate passed-in method
        self.method = method
        return self


class IPCacheItem(int):
    '''Represents ip address cache entity
    Takes one argument at creation -- ip addr as string, e.g. '127.0.0.1'
    Objects of this are used as KEYS in cache, with VALUES determining
    ban status of particular ip addr.
    '''

    def __new__(cls, ipstring):
        parts = ipstring.split('.')
        res = 0
        for i, part in enumerate(reversed(parts)):
            res |= int(part) << i*8
        return super().__new__(cls, res)

    def __eq__(self, other):
        return isinstance(other, self.__class__) and super().__eq__(other)

    def __hash__(self):
        return super().__hash__()

    def __repr__(self):
        s = '.'.join(str(int(bt)) for bt in self.to_bytes(4, 'big'))
        return "%s('%s')" % (self.__class__.__name__, s)