''' xbmcswift2.storage ~~~~~~~~~~~~~~~~~~ This module contains persistent storage classes. :copyright: (c) 2012 by Jonathan Beluch :license: GPLv3, see LICENSE for more details. ''' import os import csv import json import time try: import cPickle as pickle except ImportError: import pickle import shutil import collections from datetime import datetime from xbmcswift2.logger import log class _PersistentDictMixin(object): ''' Persistent dictionary with an API compatible with shelve and anydbm. The dict is kept in memory, so the dictionary operations run as fast as a regular dictionary. Write to disk is delayed until close or sync (similar to gdbm's fast mode). Input file format is automatically discovered. Output file format is selectable between pickle, json, and csv. All three serialization formats are backed by fast C implementations. ''' def __init__(self, filename, flag='c', mode=None, file_format='pickle'): self.flag = flag # r=readonly, c=create, or n=new self.mode = mode # None or an octal triple like 0644 self.file_format = file_format # 'csv', 'json', or 'pickle' self.filename = filename if flag != 'n' and os.access(filename, os.R_OK): log.debug('Reading %s storage from disk at "%s"', self.file_format, self.filename) fileobj = open(filename, 'rb' if file_format == 'pickle' else 'r') with fileobj: self.load(fileobj) def sync(self): '''Write the dict to disk''' if self.flag == 'r': return filename = self.filename tempname = filename + '.tmp' fileobj = open(tempname, 'wb' if self.file_format == 'pickle' else 'w') try: self.dump(fileobj) except Exception: os.remove(tempname) raise finally: fileobj.close() # shutil error (SameFileError when performing copyfile) if os.path.exists(self.filename): os.remove(self.filename) shutil.move(tempname, self.filename) # atomic commit if self.mode is not None: os.chmod(self.filename, self.mode) def close(self): '''Calls sync''' self.sync() def __enter__(self): return self def __exit__(self, *exc_info): self.close() def dump(self, fileobj): '''Handles the writing of the dict to the file object''' if self.file_format == 'csv': csv.writer(fileobj).writerows(self.raw_dict().items()) elif self.file_format == 'json': json.dump(self.raw_dict(), fileobj, separators=(',', ':')) elif self.file_format == 'pickle': pickle.dump(dict(self.raw_dict()), fileobj, 2) else: raise NotImplementedError('Unknown format: ' + repr(self.file_format)) def load(self, fileobj): '''Load the dict from the file object''' # try formats from most restrictive to least restrictive for loader in (pickle.load, json.load, csv.reader): fileobj.seek(0) try: return self.initial_update(loader(fileobj)) except Exception as e: pass raise ValueError('File not in a supported format') def raw_dict(self): '''Returns the underlying dict''' raise NotImplementedError class _Storage(collections.MutableMapping, _PersistentDictMixin): '''Storage that acts like a dict but also can persist to disk. :param filename: An absolute filepath to reprsent the storage on disk. The storage will loaded from this file if it already exists, otherwise the file will be created. :param file_format: 'pickle', 'json' or 'csv'. pickle is the default. Be aware that json and csv have limited support for python objets. .. warning:: Currently there are no limitations on the size of the storage. Please be sure to call :meth:`~xbmcswift2._Storage.clear` periodically. ''' def __init__(self, filename, file_format='pickle'): '''Acceptable formats are 'csv', 'json' and 'pickle'.''' self._items = {} _PersistentDictMixin.__init__(self, filename, file_format=file_format) def __setitem__(self, key, val): self._items.__setitem__(key, val) def __getitem__(self, key): return self._items.__getitem__(key) def __delitem__(self, key): self._items.__delitem__(key) def __iter__(self): return iter(self._items) def __len__(self): return self._items.__len__ def raw_dict(self): '''Returns the wrapped dict''' return self._items initial_update = collections.MutableMapping.update def clear(self): super(_Storage, self).clear() self.sync() class TimedStorage(_Storage): '''A dict with the ability to persist to disk and TTL for items.''' def __init__(self, filename, file_format='pickle', TTL=None): '''TTL if provided should be a datetime.timedelta. Any entries older than the provided TTL will be removed upon load and upon item access. ''' self.TTL = TTL _Storage.__init__(self, filename, file_format=file_format) def __setitem__(self, key, val, raw=False): if raw: self._items[key] = val else: self._items[key] = (val, time.time()) def __getitem__(self, key): val, timestamp = self._items[key] if self.TTL and (datetime.utcnow() - datetime.utcfromtimestamp(timestamp) > self.TTL): del self._items[key] return self._items[key][0] # Will raise KeyError return val def initial_update(self, mapping): '''Initially fills the underlying dictionary with keys, values and timestamps. ''' for key, val in mapping.items(): _, timestamp = val if not self.TTL or (datetime.utcnow() - datetime.utcfromtimestamp(timestamp) < self.TTL): self.__setitem__(key, val, raw=True)