Improve error messages for legacy and compressed formats

Author: benjeffery

python/_tskitmodule.c

@@ -1,7 +1,7 @@
 /*
  * MIT License
  *
- * Copyright (c) 2019-2022 Tskit Developers
+ * Copyright (c) 2019-2023 Tskit Developers
  * Copyright (c) 2015-2018 University of Oxford
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -221,10 +221,10 @@ handle_library_error(int err)
 {
     int kas_err;
     const char *not_kas_format_msg
-        = "File not in kastore format. If this file "
-          "was generated by msprime < 0.6.0 (June 2018) it uses the old HDF5-based "
-          "format which can no longer be read directly. Please convert to the new "
-          "kastore format using the ``tskit upgrade`` command.";
+        = "File not in kastore format. Either the file is corrupt or it is not a "
+          "tskit tree sequence file. It may be a legacy HDF file upgradable with "
+          "`tskit upgrade` or a compressed tree sequence file that can be decompressed "
+          "with `tszip`.";
     const char *ibd_pairs_not_stored_msg
         = "Sample pairs are not stored by default "
           "in the IdentitySegments object returned by ibd_segments(), and you have "

python/requirements/CI-complete/requirements.txt

@@ -19,4 +19,5 @@ pytest==7.1.3
 pytest-cov==4.0.0
 pytest-xdist==2.5.0
 svgwrite==1.4.3
+tszip==0.2.2
 xmlunittest==0.5.0

python/requirements/CI-tests-pip/requirements.txt

@@ -11,4 +11,5 @@ biopython==1.79
 dendropy==4.5.2
 networkx==2.6.3 # Held at 2.6.3 for Python 3.7 compatibility
 msgpack==1.0.4
-newick==1.3.2
+newick==1.3.2
+tszip==0.2.2

python/requirements/development.txt

@@ -36,6 +36,7 @@ sphinx-jupyterbook-latex
 sphinxcontrib-prettyspecialmethods
 tqdm
 tskit-book-theme
+tszip
 pydata_sphinx_theme>=0.7.2
 svgwrite>=1.1.10
 xmlunittest

python/tests/test_file_format.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2018-2022 Tskit Developers
+# Copyright (c) 2018-2023 Tskit Developers
 # Copyright (c) 2016-2018 University of Oxford
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -36,12 +36,12 @@
 import msprime
 import numpy as np
 import pytest
+import tszip as tszip
 
 import tests.tsutil as tsutil
 import tskit
 import tskit.exceptions as exceptions
 
-
 CURRENT_FILE_MAJOR = 12
 CURRENT_FILE_MINOR = 7
 
@@ -262,11 +262,17 @@ def test_format_too_old_raised_for_hdf5(self):
         ]
         for filename in files:
             path = os.path.join(test_data_dir, "hdf5-formats", filename)
+
             with pytest.raises(
                 exceptions.FileFormatError,
-                match="uses the old HDF5-based format which can no longer",
+                match="appears to be in HDF5 format",
             ):
                 tskit.load(path)
+            with pytest.raises(
+                exceptions.FileFormatError,
+                match="appears to be in HDF5 format",
+            ):
+                tskit.TableCollection.load(path)
 
     def test_msprime_v_0_5_0(self):
         path = os.path.join(test_data_dir, "hdf5-formats", "msprime-0.5.0_v10.0.hdf5")
@@ -511,6 +517,14 @@ def test_no_h5py(self):
             with pytest.raises(ImportError, match=msg):
                 tskit.dump_legacy(ts, path)
 
+    def test_tszip_file(self):
+        ts = msprime.simulate(5)
+        tszip.compress(ts, self.temp_file)
+        with pytest.raises(tskit.FileFormatError, match="appears to be in zip format"):
+            tskit.load(self.temp_file)
+        with pytest.raises(tskit.FileFormatError, match="appears to be in zip format"):
+            tskit.TableCollection.load(self.temp_file)
+
 
 class TestDumpFormat(TestFileFormat):
     """

python/tests/test_fileobj.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2018-2022 Tskit Developers
+# Copyright (c) 2018-2023 Tskit Developers
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -35,6 +35,7 @@
 import traceback
 
 import pytest
+import tszip
 from pytest import fixture
 
 import tskit
@@ -308,3 +309,59 @@ def verify_stream(self, ts_list, client_fd):
     def test_single_then_multi(self, ts_fixture, replicate_ts_fixture, client_fd):
         self.verify_stream([ts_fixture], client_fd)
         self.verify_stream(replicate_ts_fixture, client_fd)
+
+
+def write_to_fifo(path, file_path):
+    with open(path, "wb") as fifo:
+        with open(file_path, "rb") as file:
+            fifo.write(file.read())
+
+
+def read_from_fifo(path, expected_exception, error_text, read_func):
+    with open(path) as fifo:
+        with pytest.raises(expected_exception, match=error_text):
+            read_func(fifo)
+
+
+def write_and_read_from_fifo(fifo_path, file_path, expected_exception, error_text):
+    os.mkfifo(fifo_path)
+    for read_func in [tskit.load, tskit.TableCollection.load]:
+        read_process = multiprocessing.Process(
+            target=read_from_fifo,
+            args=(fifo_path, expected_exception, error_text, read_func),
+        )
+        read_process.start()
+        write_process = multiprocessing.Process(
+            target=write_to_fifo, args=(fifo_path, file_path)
+        )
+        write_process.start()
+        write_process.join(timeout=3)
+        read_process.join(timeout=3)
+
+
+@pytest.mark.skipif(IS_WINDOWS, reason="No FIFOs on Windows")
+class TestBadStream:
+    def test_bad_stream(self, tmp_path):
+        fifo_path = tmp_path / "fifo"
+        bad_file_path = tmp_path / "bad_file"
+        bad_file_path.write_bytes(b"bad data")
+        write_and_read_from_fifo(
+            fifo_path, bad_file_path, tskit.FileFormatError, "not in kastore format"
+        )
+
+    def test_legacy_stream(self, tmp_path):
+        fifo_path = tmp_path / "fifo"
+        legacy_file_path = os.path.join(
+            os.path.dirname(__file__), "data", "hdf5-formats", "msprime-0.3.0_v2.0.hdf5"
+        )
+        write_and_read_from_fifo(
+            fifo_path, legacy_file_path, tskit.FileFormatError, "not in kastore format"
+        )
+
+    def test_tszip_stream(self, tmp_path, ts_fixture):
+        fifo_path = tmp_path / "fifo"
+        zip_file_path = tmp_path / "tszip_file"
+        tszip.compress(ts_fixture, zip_file_path)
+        write_and_read_from_fifo(
+            fifo_path, zip_file_path, tskit.FileFormatError, "not in kastore format"
+        )

python/tskit/tables.py

@@ -3197,12 +3197,18 @@ def __getstate__(self):
     def load(cls, file_or_path, *, skip_tables=False, skip_reference_sequence=False):
         file, local_file = util.convert_file_like_to_open_file(file_or_path, "rb")
         ll_tc = _tskit.TableCollection()
-        ll_tc.load(
-            file,
-            skip_tables=skip_tables,
-            skip_reference_sequence=skip_reference_sequence,
-        )
-        return TableCollection(ll_tables=ll_tc)
+        try:
+            ll_tc.load(
+                file,
+                skip_tables=skip_tables,
+                skip_reference_sequence=skip_reference_sequence,
+            )
+            return TableCollection(ll_tables=ll_tc)
+        except tskit.FileFormatError as e:
+            util.raise_known_file_format_errors(file, e)
+        finally:
+            if local_file:
+                file.close()
 
     def dump(self, file_or_path):
         """

python/tskit/trees.py

@@ -4093,6 +4093,8 @@ def load(cls, file_or_path, *, skip_tables=False, skip_reference_sequence=False)
                 skip_reference_sequence=skip_reference_sequence,
             )
             return TreeSequence(ts)
+        except tskit.FileFormatError as e:
+            util.raise_known_file_format_errors(file, e)
         finally:
             if local_file:
                 file.close()

python/tskit/util.py

@@ -23,6 +23,7 @@
 Module responsible for various utility functions used in other modules.
 """
 import dataclasses
+import io
 import itertools
 import json
 import numbers
@@ -749,3 +750,32 @@ def random_nucleotides(length: numbers.Number, *, seed: Union[int, None] = None)
     encoded_nucleotides = np.array(list(map(ord, "ACTG")), dtype=np.int8)
     a = rng.choice(encoded_nucleotides, size=int(length))
     return a.tobytes().decode("ascii")
+
+
+def raise_known_file_format_errors(open_file, existing_exception):
+    """
+    Sniffs the file for pk-zip or hdf header bytes, then raises an exception
+    if these are detected, if not raises the existing exception.
+    """
+    # Check for HDF5 header bytes
+    try:
+        open_file.seek(0)
+        header = open_file.read(4)
+    except io.UnsupportedOperation:
+        # If we can't seek, we can't sniff the file.
+        raise existing_exception
+    if header == b"\x89HDF":
+        raise tskit.FileFormatError(
+            "The specified file appears to be in HDF5 format. This file "
+            "may have been generated by msprime < 0.6.0 (June 2018) which "
+            "can no longer be read directly. Please convert to the new "
+            "kastore format using the ``tskit upgrade`` command."
+        ) from existing_exception
+    if header[:2] == b"\x50\x4b":
+        raise tskit.FileFormatError(
+            "The specified file appears to be in zip format, so may be a compressed "
+            "tree sequence. Try using the tszip module to decompress this file before "
+            "loading. `pip install tszip; tsunzip <filename>` or use "
+            "`tszip.decompress` in Python code."
+        ) from existing_exception
+    raise existing_exception