2
0

storage.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637
  1. import cPickle as pickle
  2. import logging
  3. import os
  4. class Storage:
  5. data = None
  6. def __init__(self, storage_dir):
  7. self.logger = logging.getLogger('Storage')
  8. self.storage_dir = storage_dir
  9. self.logger.debug('Using storage at \'{0}\''.format(storage_dir))
  10. self.storage_file = open(os.path.join(storage_dir, 'storage.dat'), 'a+b')
  11. try:
  12. self.storage_file.seek(0, os.SEEK_SET)
  13. self.data = pickle.load(self.storage_file)
  14. except EOFError:
  15. self.logger.warn('The storage file was empty. I\'ll assume this is the first start with this storage directory.')
  16. except Exception as err:
  17. self.logger.error('Loading the storage failed. Please resolve the error and/or remove the storage file \'{0}\': {1}'.format(
  18. os.path.join(storage_dir, self.storage_file.name),
  19. str(err),
  20. ))
  21. raise err
  22. if self.data is None: self.data = {}
  23. self.logger.info('Opened storage with ' + str(len(self.data)) + ' entries.')
  24. def save(self):
  25. self.storage_file.seek(0, os.SEEK_SET)
  26. self.storage_file.truncate()
  27. pickle.dump(self.data, self.storage_file, protocol=2)
  28. self.storage_file.flush()
  29. def close(self):
  30. self.save()
  31. self.storage_file.close()
  32. self.data = None