Add -z option (don't compress)

This commit is contained in:
Gregory Soutade 2015-05-23 16:38:39 +02:00
parent f3fc24e45e
commit 79b58f2e1c

17
iwla.py
View File

@ -235,15 +235,18 @@ class IWLA(object):
def getDBFilename(self, time):
return os.path.join(conf.DB_ROOT, str(time.tm_year), '%02d' % (time.tm_mon), conf.DB_FILENAME)
def _openDB(self, filename, prot='r'):
if self.args.dont_compress:
return open(filename, prot)
else:
return gzip.open(filename, prot)
def _serialize(self, obj, filename):
base = os.path.dirname(filename)
if not os.path.exists(base):
os.makedirs(base)
# TODO : remove return
#return
with open(filename + '.tmp', 'wb+') as f, gzip.open(filename, 'w') as fzip:
with open(filename + '.tmp', 'wb+') as f, self._openDB(filename, 'w') as fzip:
pickle.dump(obj, f)
f.seek(0)
fzip.write(f.read())
@ -253,7 +256,7 @@ class IWLA(object):
if not os.path.exists(filename):
return None
with gzip.open(filename, 'r') as f:
with self._openDB(filename) as f:
return pickle.load(f)
return None
@ -805,6 +808,10 @@ if __name__ == '__main__':
default=False,
help='Reset analysis to a specific date (month/year)')
parser.add_argument('-z', '--dont-compress', dest='dont_compress', action='store_true',
default=False,
help='Don\'t compress databases (bigger but faster)')
args = parser.parse_args()
# Load user conf