import codecs import logging import os import pickle import shutil import tempfile import traceback import hashlib import appdirs from scandir import scandir, scandir_generic as _scandir_generic try: from collections.abc import MutableMapping unicode = str except ImportError: # Python 2 imports from collections import MutableMapping FileNotFoundError = IOError from .posixemulation import rename logger = logging.getLogger(__name__) class FileCache(MutableMapping): """A persistent file cache that is dictionary-like and has a write buffer. *appname* is passed to `appdirs `_ to determine a system-appropriate location for the cache files. The cache directory used is available via :data:`cache_dir`. By default, a write buffer is used, so writing to cache files is not done until :meth:`sync` is explicitly called. This behavior can be changed using the optional *flag* argument. .. NOTE:: Keys and values are always stored as :class:`bytes` objects. If data serialization is enabled, keys are returned as :class:`str` or :class:`unicode` objects. If data serialization is disabled, keys are returned as a :class:`bytes` object. :param str appname: The app/script the cache should be associated with. :param str flag: How the cache should be opened. See below for details. :param mode: The Unix mode for the cache files. :param str keyencoding: The encoding the keys use, defaults to 'utf-8'. This is used if *serialize* is ``False``; the keys are treated as :class:`bytes` objects. :param bool serialize: Whether or not to (de)serialize the values. If a cache is used with a :class:`~shelve.Shelf`, set this to ``False``. :param str app_cache_dir: absolute path to root cache directory to be used in place of system-appropriate location determined by appdirs The optional *flag* argument can be: +---------+-------------------------------------------+ | Value | Meaning | +=========+===========================================+ | ``'r'`` | Open existing cache for reading only | +---------+-------------------------------------------+ | ``'w'`` | Open existing cache for reading and | | | writing | +---------+-------------------------------------------+ | ``'c'`` | Open cache for reading and writing, | | | creating it if it doesn't exist (default) | +---------+-------------------------------------------+ | ``'n'`` | Always create a new, empty cache, open | | | for reading and writing | +---------+-------------------------------------------+ If a ``'s'`` is appended to the *flag* argument, the cache will be opened in sync mode. Writing to the cache will happen immediately and will not be buffered. If an application needs to use more than one cache, then it should use subcaches. To create a subcache, append a series of one or more names separated by periods to the application name when creating a :class:`FileCache` object (e.g. ``'appname.subcache'`` or ``'appname.subcache.subcache'``). Subcaches are a way for an application to use more than one cache without polluting a user's cache directory. All caches -- main caches or subcaches -- are totally independent. The only aspect in which they are linked is that all of an application's caches exist in the same system directory. Because each cache is independent of every other cache, calling :meth:`delete` on an application's main cache will not delete data in its subcaches. """ def __init__(self, appname, flag='c', mode=0o666, keyencoding='utf-8', serialize=True, app_cache_dir=None, key_file_ext=".txt"): """Initialize a :class:`FileCache` object.""" if not isinstance(flag, str): raise TypeError("flag must be str not '{}'".format(type(flag))) elif flag[0] not in 'rwcn': raise ValueError("invalid flag: '{}', first flag must be one of " "'r', 'w', 'c' or 'n'".format(flag)) elif len(flag) > 1 and flag[1] != 's': raise ValueError("invalid flag: '{}', second flag must be " "'s'".format(flag)) appname, subcache = self._parse_appname(appname) if 'cache' in subcache: raise ValueError("invalid subcache name: 'cache'.") self._is_subcache = bool(subcache) if not app_cache_dir: app_cache_dir = appdirs.user_cache_dir(appname, appname) subcache_dir = os.path.join(app_cache_dir, *subcache) self.cache_dir = os.path.join(subcache_dir, 'cache') exists = os.path.exists(self.cache_dir) if len(flag) > 1 and flag[1] == 's': self._sync = True else: self._sync = False self._buffer = {} if exists and 'n' in flag: self.clear() self.create() elif not exists and ('c' in flag or 'n' in flag): self.create() elif not exists: raise FileNotFoundError("no such directory: '{}'".format( self.cache_dir)) self._flag = 'rb' if 'r' in flag else 'wb' self._mode = mode self._keyencoding = keyencoding self._serialize = serialize self.key_file_ext = key_file_ext def _parse_appname(self, appname): """Splits an appname into the appname and subcache components.""" components = appname.split('.') return components[0], components[1:] def create(self): """Create the write buffer and cache directory.""" if not self._sync and not hasattr(self, '_buffer'): self._buffer = {} if not os.path.exists(self.cache_dir): os.makedirs(self.cache_dir) def clear(self): """Remove all items from the write buffer and cache. The write buffer object and cache directory are not deleted. """ self.delete() self.create() def delete(self): """Delete the write buffer and cache directory.""" if not self._sync: del self._buffer shutil.rmtree(self.cache_dir) def close(self): """Sync the write buffer, then close the cache. If a closed :class:`FileCache` object's methods are called, a :exc:`ValueError` will be raised. """ self.sync() self.sync = self.create = self.delete = self._closed self._write_to_file = self._read_to_file = self._closed self._key_to_filename = self._filename_to_key = self._closed self.__getitem__ = self.__setitem__ = self.__delitem__ = self._closed self.__iter__ = self.__len__ = self.__contains__ = self._closed def sync(self): """Sync the write buffer with the cache files and clear the buffer. If the :class:`FileCache` object was opened with the optional ``'s'`` *flag* argument, then calling :meth:`sync` will do nothing. """ if self._sync: return # opened in sync mode, so skip the manual sync self._sync = True for ekey in self._buffer: filename = self._key_to_filename(ekey) try: self._write_to_file(filename, self._buffer[ekey]) except: logger.error("Couldn't write content from %r to cache file: %r: %s", ekey, filename, traceback.format_exc()) try: self.__write_to_file(filename + self.key_file_ext, ekey) except: logger.error("Couldn't write content from %r to cache file: %r: %s", ekey, filename, traceback.format_exc()) self._buffer.clear() self._sync = False def _closed(self, *args, **kwargs): """Filler method for closed cache methods.""" raise ValueError("invalid operation on closed cache") def _encode_key(self, key): """ Keys are implicitly converted to :class:`bytes` if passed as :class:`str`. """ if isinstance(key, str) or isinstance(key, unicode): key = key.encode(self._keyencoding) elif not isinstance(key, bytes): raise TypeError("key must be bytes or str") return key.decode(self._keyencoding) def _decode_key(self, key): """ Keys are returned as :class:`str` if serialization is enabled. Keys are returned as :class:`bytes` if serialization is disabled. """ bkey = key.encode(self._keyencoding) return bkey.decode(self._keyencoding) if self._serialize else bkey def _dumps(self, value): return value if not self._serialize else pickle.dumps(value) def _loads(self, value): return value if not self._serialize else pickle.loads(value) def _key_to_filename(self, key): """Convert an encoded key to an absolute cache filename.""" if isinstance(key, unicode): key = key.encode(self._keyencoding) return os.path.join(self.cache_dir, hashlib.md5(key).hexdigest()) def _filename_to_key(self, absfilename): """Convert an absolute cache filename to a key name.""" hkey_hdr_fn = absfilename + self.key_file_ext if os.path.isfile(hkey_hdr_fn): with open(hkey_hdr_fn, 'rb') as f: key = f.read() return key.decode(self._keyencoding) if self._serialize else key def _all_filenames(self, scandir_generic=True): """Return a list of absolute cache filenames""" _scandir = _scandir_generic if scandir_generic else scandir try: for entry in _scandir(self.cache_dir): if entry.is_file(follow_symlinks=False) and not entry.name.endswith(self.key_file_ext): yield os.path.join(self.cache_dir, entry.name) except (FileNotFoundError, OSError): raise StopIteration def _all_keys(self): """Return a list of all encoded key names.""" file_keys = [self._filename_to_key(fn) for fn in self._all_filenames()] if self._sync: return set(file_keys) else: return set(file_keys + list(self._buffer)) def __write_to_file(self, filename, value): """Write bytesvalue to filename.""" fh, tmp = tempfile.mkstemp() with os.fdopen(fh, self._flag) as f: f.write(value) rename(tmp, filename) os.chmod(filename, self._mode) def _write_to_file(self, filename, bytesvalue): self.__write_to_file(filename, self._dumps(bytesvalue)) def _read_from_file(self, filename): """Read data from filename.""" try: with open(filename, 'rb') as f: return self._loads(f.read()) except (IOError, OSError): logger.warning('Error opening file: {}'.format(filename)) return None def __setitem__(self, key, value): ekey = self._encode_key(key) if not self._sync: self._buffer[ekey] = value else: filename = self._key_to_filename(ekey) self._write_to_file(filename, value) self.__write_to_file(filename + self.key_file_ext, ekey) def __getitem__(self, key): ekey = self._encode_key(key) if not self._sync: try: return self._buffer[ekey] except KeyError: pass filename = self._key_to_filename(ekey) if not os.path.isfile(filename): raise KeyError(key) return self._read_from_file(filename) def __delitem__(self, key): ekey = self._encode_key(key) filename = self._key_to_filename(ekey) if not self._sync: try: del self._buffer[ekey] except KeyError: if filename not in self._all_filenames(): raise KeyError(key) try: os.remove(filename) except (IOError, OSError): pass try: os.remove(filename + self.key_file_ext) except (IOError, OSError): pass def __iter__(self): for key in self._all_keys(): yield self._decode_key(key) def __len__(self): return len(self._all_keys()) def __contains__(self, key): ekey = self._encode_key(key) if not self._sync: try: return ekey in self._buffer except KeyError: pass filename = self._key_to_filename(ekey) return os.path.isfile(filename)