cache.py 4.6 KB
Newer Older
1 2 3 4
# Slixmpp: The Slick XMPP Library
# Copyright (C) 2018 Emmanuel Gil Peyrot
# This file is part of Slixmpp.
# See the file LICENSE for copying permission.
5 6 7
import os
import logging

mathieui's avatar
mathieui committed
8 9
from typing import Callable, Optional, Any

10 11
log = logging.getLogger(__name__)

mathieui's avatar
mathieui committed
12

13 14 15 16 17 18 19
class Cache:
    def retrieve(self, key):
        raise NotImplementedError

    def store(self, key, value):
        raise NotImplementedError

20
    def remove(self, key):
mathieui's avatar
mathieui committed
21 22
        raise NotImplementedError

23

24 25 26 27 28 29 30
class PerJidCache:
    def retrieve_by_jid(self, jid, key):
        raise NotImplementedError

    def store_by_jid(self, jid, key, value):
        raise NotImplementedError

31 32 33
    def remove_by_jid(self, jid, key):
        raise NotImplementedError

mathieui's avatar
mathieui committed
34

35 36 37 38 39 40 41 42 43 44 45
class MemoryCache(Cache):
    def __init__(self):
        self.cache = {}

    def retrieve(self, key):
        return self.cache.get(key, None)

    def store(self, key, value):
        self.cache[key] = value
        return True

46 47 48 49 50
    def remove(self, key):
        if key in self.cache:
            del self.cache[key]
        return True

mathieui's avatar
mathieui committed
51

52 53 54 55 56
class MemoryPerJidCache(PerJidCache):
    def __init__(self):
        self.cache = {}

    def retrieve_by_jid(self, jid, key):
57
        cache = self.cache.get(jid, None)
58 59 60 61 62
        if cache is None:
            return None
        return cache.get(key, None)

    def store_by_jid(self, jid, key, value):
63
        cache = self.cache.setdefault(jid, {})
64 65 66
        cache[key] = value
        return True

67 68 69 70 71 72
    def remove_by_jid(self, jid, key):
        cache = self.cache.get(jid, None)
        if cache is not None and key in cache:
            del cache[key]
        return True

mathieui's avatar
mathieui committed
73

74
class FileSystemStorage:
mathieui's avatar
mathieui committed
75
    def __init__(self, encode: Optional[Callable[[Any], str]], decode: Optional[Callable[[str], Any]], binary: bool):
76 77
        self.encode = encode if encode is not None else lambda x: x
        self.decode = decode if decode is not None else lambda x: x
78 79 80
        self.read = 'rb' if binary else 'r'
        self.write = 'wb' if binary else 'w'

mathieui's avatar
mathieui committed
81
    def _retrieve(self, directory: str, key: str):
82 83 84 85 86 87 88 89
        filename = os.path.join(directory, key.replace('/', '_'))
        try:
            with open(filename, self.read) as cache_file:
                return self.decode(cache_file.read())
        except FileNotFoundError:
            log.debug('%s not present in cache', key)
        except OSError:
            log.debug('Failed to read %s from cache:', key, exc_info=True)
90 91 92 93
        except Exception:
            log.debug('Failed to decode %s from cache:', key, exc_info=True)
            log.debug('Removing %s entry', key)
            self._remove(directory, key)
94

mathieui's avatar
mathieui committed
95
    def _store(self, directory: str, key: str, value):
96 97 98 99 100 101 102 103 104
        filename = os.path.join(directory, key.replace('/', '_'))
        try:
            os.makedirs(directory, exist_ok=True)
            with open(filename, self.write) as output:
                output.write(self.encode(value))
                return True
        except OSError:
            log.debug('Failed to store %s to cache:', key, exc_info=True)
            return False
105 106
        except Exception:
            log.debug('Failed to encode %s to cache:', key, exc_info=True)
107

mathieui's avatar
mathieui committed
108
    def _remove(self, directory: str, key: str):
109 110 111 112 113 114 115 116
        filename = os.path.join(directory, key.replace('/', '_'))
        try:
            os.remove(filename)
        except OSError:
            log.debug('Failed to remove %s from cache:', key, exc_info=True)
            return False
        return True

mathieui's avatar
mathieui committed
117

118
class FileSystemCache(Cache, FileSystemStorage):
mathieui's avatar
mathieui committed
119
    def __init__(self, directory: str, cache_type: str, *, encode=None, decode=None, binary=False):
120 121 122 123 124 125 126 127 128
        FileSystemStorage.__init__(self, encode, decode, binary)
        self.base_dir = os.path.join(directory, cache_type)

    def retrieve(self, key):
        return self._retrieve(self.base_dir, key)

    def store(self, key, value):
        return self._store(self.base_dir, key, value)

129 130 131
    def remove(self, key):
        return self._remove(self.base_dir, key)

mathieui's avatar
mathieui committed
132

133
class FileSystemPerJidCache(PerJidCache, FileSystemStorage):
mathieui's avatar
mathieui committed
134
    def __init__(self, directory: str, cache_type: str, *, encode=None, decode=None, binary=False):
135 136 137 138
        FileSystemStorage.__init__(self, encode, decode, binary)
        self.base_dir = os.path.join(directory, cache_type)

    def retrieve_by_jid(self, jid, key):
139
        directory = os.path.join(self.base_dir, jid)
140 141 142
        return self._retrieve(directory, key)

    def store_by_jid(self, jid, key, value):
143
        directory = os.path.join(self.base_dir, jid)
144
        return self._store(directory, key, value)
145 146 147 148

    def remove_by_jid(self, jid, key):
        directory = os.path.join(self.base_dir, jid)
        return self._remove(directory, key)