BATCAVE/ffstatus/storage.py

import cPickle as pickle
import logging
import os

class Storage:
	data = None

	def __init__(self, storage_dir):
		self.logger = logging.getLogger('Storage')
		self.storage_dir = storage_dir

		self.logger.debug('Using storage at \'{0}\''.format(storage_dir))
		self.storage_file = open(os.path.join(storage_dir, 'storage.dat'), 'a+b')
		try:
			self.storage_file.seek(0, os.SEEK_SET)
			self.data = pickle.load(self.storage_file)
		except EOFError:
			self.logger.warn('The storage file was empty. I\'ll assume this is the first start with this storage directory.')
		except Exception as err:
			self.logger.error('Loading the storage failed. Please resolve the error and/or remove the storage file \'{0}\': {1}'.format(
				os.path.join(storage_dir, self.storage_file.name),
				str(err),
			))
			raise err
		if self.data is None: self.data = {}
		self.logger.info('Opened storage with ' + str(len(self.data)) + ' entries.')

	def save(self):
		self.storage_file.seek(0, os.SEEK_SET)
		self.storage_file.truncate()
		pickle.dump(self.data, self.storage_file, protocol=2)
		self.storage_file.flush()

	def close(self):
		self.save()
		self.storage_file.close()
		self.data = None