soutade/iwla
1
0
Fork 0
Code Issues Pull Requests Releases 2 Activity

Add -z option (don't compress)

Browse Source
This commit is contained in:
Gregory Soutade 2015-05-23 16:38:39 +02:00
parent f3fc24e45e
commit 79b58f2e1c
1 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
iwla.py
View File

@ -235,15 +235,18 @@ class IWLA(object):
def getDBFilename(self, time): def getDBFilename(self, time):
return os.path.join(conf.DB_ROOT, str(time.tm_year), '%02d' % (time.tm_mon), conf.DB_FILENAME) 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): def _serialize(self, obj, filename):
base = os.path.dirname(filename) base = os.path.dirname(filename)
if not os.path.exists(base): if not os.path.exists(base):
os.makedirs(base) os.makedirs(base)
# TODO : remove return with open(filename + '.tmp', 'wb+') as f, self._openDB(filename, 'w') as fzip:
#return
with open(filename + '.tmp', 'wb+') as f, gzip.open(filename, 'w') as fzip:
pickle.dump(obj, f) pickle.dump(obj, f)
f.seek(0) f.seek(0)
fzip.write(f.read()) fzip.write(f.read())
@ -253,7 +256,7 @@ class IWLA(object):
if not os.path.exists(filename): if not os.path.exists(filename):
return None return None
with gzip.open(filename, 'r') as f: with self._openDB(filename) as f:
return pickle.load(f) return pickle.load(f)
return None return None
@ -805,6 +808,10 @@ if __name__ == '__main__':
default=False, default=False,
help='Reset analysis to a specific date (month/year)') 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() args = parser.parse_args()
# Load user conf # Load user conf