read compressed log files according to file ext

cublas_parse.py

@@ -2,6 +2,7 @@
 
 from tabulate import tabulate
 from mvt.display import cuBLASApothecary
+from mvt.open import open_compressed
 from tqdm import tqdm
 
 def parse_iter_blask(f, time_thr = 1e-6):
@@ -60,7 +61,7 @@ def parse_and_display(f,count):
 
     args = parser.parse_args()
     if args.filename:
-        f = open(args.filename, 'r')
+        f = open_compressed(args.filename, 'rt')
     elif not sys.stdin.isatty():
         f = sys.stdin 
     else:

mkl_parse.py

@@ -3,6 +3,7 @@
 from tabulate import tabulate
 from mvt.display import BLASApothecary, FFTApothecary
 from mvt.parse import parse_iter
+from mvt.open import open_compressed
 import greenlet
 from tqdm import tqdm
 
@@ -89,7 +90,7 @@ def parse_and_display(f,count):
 
     args = parser.parse_args()
     if args.filename:
-        f = open(args.filename, 'r')
+        f = open_compressed(args.filename, 'rt')
     elif not sys.stdin.isatty():
         f = sys.stdin
     else:

mvt/open.py

@@ -0,0 +1,18 @@
+"""
+File open helpers.
+"""
+
+def open_compressed(file_path, mode='r'):
+    """Open file path and decompress on the fly using built in standard
+    library modules if it has the appropriate extension."""
+    if file_path.endswith('.xz'):
+        import lzma
+        return lzma.open(file_path, mode)
+    elif file_path.endswith('.bz2'):
+        import bz2
+        return bz2.open(file_path, mode)
+    elif file_path.endswith('.gz'):
+        import gzip
+        return gzip.open(file_path, mode)
+    else:
+        return open(file_path, mode)