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 class _PersistentDictMixin(object): 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): 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() try: os.remove(self.filename) except: pass try: shutil.copyfile(tempname, self.filename) # atomic commit except: raise 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)