Support for reading and writing compressed dipha files.

3 files changed, 18 insertions, 10 deletions

diff --git a/phstuff/diphawrapper.py b/phstuff/diphawrapper.py
index f0a61de..930fbe2 100644
--- a/phstuff/diphawrapper.py
+++ b/phstuff/diphawrapper.py
@@ -9,6 +9,7 @@ import subprocess
 
 import phstuff.barcode as bc
 import phstuff.simplicial as simpl
+import phstuff.misc as misc
 
 """Handle marshalling data into and out of DIPHA's formats, and
 optionally manage running the DIPHA executable.
@@ -57,7 +58,7 @@ def save_weight_matrix(fname, weights):
     if weights.shape[0] != weights.shape[1]:
         raise ValueError("Matrix is not square.")
 
-    with open(fname, "wb") as f:
+    with misc.gz_wrapper(fname, "wb") as f:
         f.write(struct.pack("<q", DIPHA_MAGIC))
         f.write(struct.pack("<q", DIPHA_DISTANCE_MATRIX))
         f.write(struct.pack("<q", m))
@@ -85,7 +86,7 @@ def load_weight_matrix(fname):
 
     """
 
-    with open(fname, "rb") as f:
+    with misc.gz_wrapper(fname, "rb") as f:
         if struct.unpack('<q', f.read(8))[0] != DIPHA_MAGIC:
             raise IOError("File %s is not a valid DIPHA file." %(fname))
         if struct.unpack('<q', f.read(8))[0] != DIPHA_DISTANCE_MATRIX:
@@ -130,7 +131,7 @@ def save_masked_weight_matrix(fname, weights):
 
     n = weights.shape[1]
 
-    with open(fname, "wb") as f:
+    with misc.gz_wrapper(fname, "wb") as f:
         f.write(struct.pack("<q", DIPHA_MAGIC))
         f.write(struct.pack("<q", DIPHA_SPARSE_DISTANCE_MATRIX))
         f.write(struct.pack("<q", n))
@@ -181,7 +182,7 @@ def save_edge_list(fname, edge_list):
     
     n = len(edge_list)
 
-    with open(fname, "wb") as f:
+    with misc.gz_wrapper(fname, "wb") as f:
         f.write(struct.pack("<q", DIPHA_MAGIC))
         f.write(struct.pack("<q", DIPHA_SPARSE_DISTANCE_MATRIX))
         f.write(struct.pack("<q", n))
@@ -212,7 +213,7 @@ def save_cubical(fname, array):
 
     """
 
-    with open(fname, "wb") as f:
+    with misc.gz_wrapper(fname, "wb") as f:
         f.write(struct.pack("<q", DIPHA_MAGIC))
         f.write(struct.pack("<q", DIPHA_IMAGE_DATA))
         f.write(struct.pack("<q", np.product(array.shape)))
@@ -242,7 +243,7 @@ def save_simplicial(fname, complex):
     complex.order()
     assert(complex.is_ordered())
     
-    with open(fname, "wb") as f:
+    with misc.gz_wrapper(fname, "wb") as f:
         f.write(struct.pack("<q", DIPHA_MAGIC))
         f.write(struct.pack("<q", DIPHA_WEIGHTED_BOUNDARY_MATRIX))
         f.write(struct.pack("<q", 0))
@@ -294,7 +295,7 @@ def load_barcode(fname, top_dim = None):
     
     ret = dict()
     
-    with open(fname, "rb") as f:
+    with misc.gz_wrapper(fname, "rb") as f:
         if struct.unpack('<q', f.read(8))[0] != DIPHA_MAGIC:
             raise IOError("File %s is not a valid DIPHA file." %(fname))
         if struct.unpack('<q', f.read(8))[0] != DIPHA_PERSISTENCE_DIAGRAM:
@@ -331,7 +332,7 @@ def save_text_barcode(fname, barcode):
 
     """
 
-    with open(fname, "w") as f:
+    with misc.gz_wrapper(fname, "wt") as f:
         for dim in barcode.keys():
             for interval in barcode[dim]:
                 if interval.is_finite():
@@ -352,7 +353,7 @@ def save_barcode(fname, barcode):
 
     """
 
-    with open(fname, "wb") as f:
+    with misc.gz_wrapper(fname, "wb") as f:
         f.write(struct.pack("<q", DIPHA_MAGIC))
         f.write(struct.pack("<q", DIPHA_PERSISTENCE_DIAGRAM))
 
diff --git a/phstuff/misc.py b/phstuff/misc.py
new file mode 100644
index 0000000..76046fe
--- /dev/null
+++ b/phstuff/misc.py
@@ -0,0 +1,7 @@
+import gzip
+
+def gz_wrapper(fname, mode):
+    if fname.endswith(".gz"):
+        return gzip.open(fname, mode)
+    else:
+        return open(fname, mode)
diff --git a/setup.py b/setup.py
index fbc4ba3..696b9d5 100644
--- a/setup.py
+++ b/setup.py
@@ -1,7 +1,7 @@
 from setuptools import setup
 
 setup(name="phstuff",
-      version="0.0.10", # Remember to update version and download
+      version="0.0.11", # Remember to update version and download
                         # information in README.md too!
       description="Some persistent homology glue",
       url="https://nonempty.org/software/python-phstuff",