sima/lib/cache.py

   1 # -*- coding: utf-8 -*-
   2
   3 # Copyright (c) 2014, 2021 kaliko <kaliko@azylum.org>
   4 # Copyright (c) 2012, 2013 Eric Larson <eric@ionrock.org>
   5 #
   6 #   This program is free software: you can redistribute it and/or modify
   7 #   it under the terms of the GNU General Public License as published by
   8 #   the Free Software Foundation, either version 3 of the License, or
   9 #   (at your option) any later version.
  10 #
  11 #   This program is distributed in the hope that it will be useful,
  12 #   but WITHOUT ANY WARRANTY; without even the implied warranty of
  13 #   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14 #   GNU General Public License for more details.
  15 #
  16 #   You should have received a copy of the GNU General Public License
  17 #   along with this program.  If not, see <http://www.gnu.org/licenses/>.
  18 #
  19 #
  20 """
  21 The cache object API for implementing caches. The default is just a
  22 dictionary, which in turns means it is not threadsafe for writing.
  23 """
  24
  25 import os
  26 import codecs
  27
  28 from hashlib import md5
  29 from pickle import load, dump
  30 from threading import Lock
  31
  32 from ..utils.filelock import FileLock
  33
  34
  35 class BaseCache:
  36
  37     def get(self, key):
  38         """Get cache value"""
  39         raise NotImplementedError
  40
  41     def set(self, key, value):
  42         """Set cache value"""
  43         raise NotImplementedError
  44
  45     def delete(self, key):
  46         """Remove cache value"""
  47         raise NotImplementedError
  48
  49
  50 class DictCache(BaseCache):
  51
  52     def __init__(self, init_dict=None):
  53         self.lock = Lock()
  54         self.data = init_dict or {}
  55
  56     def get(self, key):
  57         return self.data.get(key, None)
  58
  59     def set(self, key, value):
  60         with self.lock:
  61             self.data.update({key: value})
  62
  63     def delete(self, key):
  64         with self.lock:
  65             if key in self.data:
  66                 self.data.pop(key)
  67
  68
  69 class FileCache:
  70
  71     def __init__(self, directory, forever=False):
  72         self.directory = directory
  73         self.forever = forever
  74
  75         if not os.path.isdir(self.directory):
  76             os.makedirs(self.directory, mode=0o755)
  77
  78     def encode(self, val):
  79         return md5(val.encode('utf-8')).hexdigest()
  80
  81     def _fn(self, name):
  82         return os.path.join(self.directory, self.encode(name))
  83
  84     def get(self, key):
  85         name = self._fn(key)
  86         if os.path.exists(name):
  87             return load(codecs.open(name, 'rb'))
  88         return None
  89
  90     def set(self, key, value):
  91         name = self._fn(key)
  92         with FileLock(name):
  93             with codecs.open(name, 'w+b') as flh:
  94                 dump(value, flh)
  95
  96     def delete(self, key):
  97         if not self.forever:
  98             os.remove(self._fn(key))
  99
 100     def __iter__(self):
 101         for dirpath, _, filenames in os.walk(self.directory):
 102             for item in filenames:
 103                 name = os.path.join(dirpath, item)
 104                 yield load(codecs.open(name, 'rb'))