Repository: tskit-dev/tszip

Issue:
Load from compressed or uncompressed 
It would be very handy to have a function that loads either an uncompressed tskit file or a tszipped one. 

I guess mimicking tskit.load would be a natural choice?

Discussion:

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 7f5a4a9..86c28c3 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,3 +1,11 @@
+--------------------
+[0.2.3] - 2023-XX-XX
+--------------------
+
+- Add `tszip.load` which loads both compressed and uncompressed trees sequences
+  (benjeffery, #75)
+
+
 --------------------
 [0.2.2] - 2022-02-22
 --------------------
diff --git a/tszip/__init__.py b/tszip/__init__.py
index 5c6f4ea..9d8feec 100644
--- a/tszip/__init__.py
+++ b/tszip/__init__.py
@@ -21,5 +21,6 @@
 # SOFTWARE.
 from .compression import compress  # NOQA
 from .compression import decompress  # NOQA
+from .compression import load  # NOQA
 from .compression import print_summary  # NOQA
 from .provenance import __version__  # NOQA
diff --git a/tszip/compression.py b/tszip/compression.py
index e23e060..6bea6fd 100644
--- a/tszip/compression.py
+++ b/tszip/compression.py
@@ -388,3 +388,34 @@ def print_summary(path, verbosity=0):
         if verbosity > 0:
             for line in str(array.info).splitlines():
                 print("\t", line)
+
+
+def load(path):
+    """
+    Open a tszip or normal tskit file. This is a convenience function that
+    determines if the file needs to be decompressed or not, returning
+    the tree sequence instance in either case.
+
+    :param str path: The location of the tszip compressed file or
+        standard tskit file to load.
+    :rtype: tskit.TreeSequence
+    :return: A :class:`tskit.TreeSequence` instance corresponding to
+        the specified file.
+    """
+    path = str(path)
+
+    # Determine if the file is a zip file, this seems more robust than
+    # checking the file extension, or depending on exceptions. Note that
+    # `is_zipfile` not only checks the header but also the EOCD record at
+    # then end of the file. This means we read the file twice, but as
+    # tree sequences are usually less than a few GB this should not
+    # be a problem.
+    with open(path, "rb") as f:
+        is_zip = zipfile.is_zipfile(f)
+    if is_zip:
+        return decompress(path)
+    else:
+        # Open everything else with tskit. We could check for a
+        # kastore header here, but this way we get all the normal
+        # tskit exceptions on error
+        return tskit.load(path)


Test patch:
diff --git a/tests/test_compression.py b/tests/test_compression.py
index b336074..e46ccfa 100644
--- a/tests/test_compression.py
+++ b/tests/test_compression.py
@@ -28,6 +28,7 @@ import unittest
 
 import msprime
 import numpy as np
+import pytest
 import tskit
 import zarr
 
@@ -383,3 +384,32 @@ class TestFileErrors(unittest.TestCase):
         ts = msprime.simulate(10, random_seed=1)
         with self.assertRaises(OSError):
             tszip.compress(ts, self.path.parent)
+
+
+class TestLoad:
+    def test_missing_file(self):
+        path = "/no/such/file"
+        with pytest.raises(FileNotFoundError):
+            tszip.load(path)
+
+    def test_load_dir(self):
+        with pytest.raises(OSError):
+            tszip.load(pathlib.Path(__file__).parent)
+
+    def test_wrong_format(self, tmpdir):
+        path = pathlib.Path(tmpdir) / "treeseq.tsz"
+        with open(str(path), "w") as f:
+            f.write("")
+        with pytest.raises(EOFError):
+            tszip.load(path)
+        for contents in ["1234", "X" * 1024]:
+            with open(str(path), "w") as f:
+                f.write(contents)
+            with pytest.raises(tskit.FileFormatError):
+                tszip.load(path)
+
+    def test_open_both(self):
+        files = pathlib.Path(__file__).parent / "files"
+        ts = tszip.load(files / "1.0.0.trees.tsz")
+        ts2 = tszip.load(files / "1.0.0.trees")
+        assert ts == ts2