parent 340315a

author Arjun Biddanda <[email protected]> 1636730041 -0500
committer Arjun Biddanda <[email protected]> 1645130969 -0500

parent 340315a
author Arjun Biddanda <[email protected]> 1636730041 -0500
committer Arjun Biddanda <[email protected]> 1645130930 -0500

feature: tszip compression to stdout

Author: aabiddanda

tests/test_cli.py

@@ -36,6 +36,17 @@
 import tszip.cli as cli
 
 
+def get_stdout_for_pytest():
+    """
+    Pytest automatically wraps stdout to intercept output, but the object
+    that it uses isn't fully compatible with the production implementation.
+    Specifically, it doesn't provide a "buffer" attribute, which we
+    need when writing binary data to it. This is a workaround to make
+    out tests work.
+    """
+    return sys.stdout
+
+
 class TestException(Exception):
     """
     Custom exception we can throw for testing.
@@ -284,15 +295,22 @@ def test_bad_file_format(self):
                 f" command."
             )
 
-    def test_compress_stdout(self):
+    def test_compress_stdout_keep(self):
         self.assertTrue(self.trees_path.exists())
-        with mock.patch("tszip.cli.exit", side_effect=TestException) as mocked_exit:
-            with self.assertRaises(TestException):
-                self.run_tszip([str(self.trees_path)] + ["-c"])
-            mocked_exit.assert_called_once_with(
-                "Compressing to stdout not currently supported;"
-                "Please see https://github.com/tskit-dev/tszip/issues/49"
-            )
+        with mock.patch("tszip.cli.get_stdout", wraps=get_stdout_for_pytest):
+            self.run_tszip_stdout([str(self.trees_path)] + ["-c"])
+        self.assertTrue(self.trees_path.exists())
+
+    def test_compress_stdout_correct(self):
+        self.assertTrue(self.trees_path.exists())
+        tmp_file = pathlib.Path(self.tmpdir.name) / "stdout.trees"
+        with mock.patch("tszip.cli.get_stdout", wraps=get_stdout_for_pytest):
+            stdout, stderr = self.run_tszip_stdout(["-c", str(self.trees_path)])
+        with open(tmp_file, "wb+") as tmp:
+            tmp.write(stdout)
+        self.assertTrue(tmp_file.exists())
+        ts = tszip.decompress(str(tmp_file))
+        self.assertEqual(ts.tables, self.ts.tables)
 
 
 class DecompressSemanticsMixin:
@@ -416,7 +434,7 @@ def test_bad_file_format(self):
             with self.assertRaises(TestException):
                 self.run_decompress([str(self.compressed_path)])
             mocked_exit.assert_called_once_with(
-                "Error reading '{}': File is not in tgzip format".format(
+                "Error reading '{}': File is not in tszip format".format(
                     self.compressed_path
                 )
             )
@@ -485,7 +503,7 @@ def test_bad_file_format(self):
             with self.assertRaises(TestException):
                 cli.tszip_main([str(self.compressed_path), "-l"])
             mocked_exit.assert_called_once_with(
-                "Error reading '{}': File is not in tgzip format".format(
+                "Error reading '{}': File is not in tszip format".format(
                     self.compressed_path
                 )
             )

tszip/cli.py

@@ -98,16 +98,18 @@ def remove_input(infile, args):
 
 def check_output(outfile, args):
     if outfile.exists():
-        if not args.force:
+        if not (args.force or args.stdout):
             exit(f"'{outfile}' already exists; use --force to overwrite")
 
 
+# Allows us to easily patch when running tests.
+def get_stdout():
+    return sys.stdout.buffer
+
+
 def run_compress(args):
     if args.stdout:
-        exit(
-            "Compressing to stdout not currently supported;"
-            "Please see https://github.com/tskit-dev/tszip/issues/49"
-        )
+        args.keep = True
     setup_logging(args)
     for file_arg in args.files:
         logger.info(f"Compressing {file_arg}")
@@ -119,6 +121,8 @@ def run_compress(args):
         infile = pathlib.Path(file_arg)
         outfile = pathlib.Path(file_arg + args.suffix)
         check_output(outfile, args)
+        if args.stdout:
+            outfile = get_stdout()
         tszip.compress(ts, outfile, variants_only=args.variants_only)
         remove_input(infile, args)
 

tszip/compression.py

@@ -23,9 +23,11 @@
 Compression utilities for tskit tree sequences.
 """
 import contextlib
+import functools
 import json
 import logging
-import os.path
+import os
+import pathlib
 import tempfile
 import warnings
 import zipfile
@@ -71,32 +73,43 @@ def minimal_dtype(array):
     return dtype
 
 
-def compress(ts, path, variants_only=False):
+def compress(ts, destination, variants_only=False):
     """
-    Compresses the specified tree sequence and writes it to the specified path.
-    By default, fully lossless compression is used so that tree sequences are
-    identical before and after compression. By specifying the ``variants_only``
-    option, a lossy compression can be used, which discards any information
-    that is not needed to represent the variants (which are stored losslessly).
+    Compresses the specified tree sequence and writes it to the specified path
+    or file-like object. By default, fully lossless compression is used so that
+    tree sequences are identical before and after compression. By specifying
+    the ``variants_only`` option, a lossy compression can be used, which
+    discards any information that is not needed to represent the variants
+    (which are stored losslessly).
 
     :param tskit.TreeSequence ts: The input tree sequence.
-    :param str destination: The string or :class:`pathlib.Path` instance describing
-        the location of the compressed file.
+    :param str destination: The string, :class:`pathlib.Path` or file-like object
+        we should write the compressed file to.
     :param bool variants_only: If True, discard all information not necessary
         to represent the variants in the input file.
     """
-    destination = str(path)
-    # Write the file into a temporary directory on the same file system so that
-    # we can write the output atomically.
-    destdir = os.path.dirname(os.path.abspath(destination))
+    try:
+        destination = pathlib.Path(destination).resolve()
+        is_path = True
+        destdir = destination.parent
+    except TypeError:
+        is_path = False
+        destdir = None
     with tempfile.TemporaryDirectory(dir=destdir, prefix=".tszip_work_") as tmpdir:
-        filename = os.path.join(tmpdir, "tmp.trees.tgz")
+        filename = pathlib.Path(tmpdir, "tmp.trees.tgz")
         logging.debug(f"Writing to temporary file {filename}")
         with zarr.ZipStore(filename, mode="w") as store:
             root = zarr.group(store=store)
             compress_zarr(ts, root, variants_only=variants_only)
-        os.replace(filename, destination)
-    logging.info(f"Wrote {destination}")
+        if is_path:
+            os.replace(filename, destination)
+            logging.info(f"Wrote {destination}")
+        else:
+            # Assume that destination is a file-like object open in "wb" mode.
+            with open(filename, "rb") as source:
+                chunk_size = 2 ** 10  # 1MiB
+                for chunk in iter(functools.partial(source.read, chunk_size), b""):
+                    destination.write(chunk)
 
 
 def decompress(path):
@@ -275,7 +288,7 @@ def load_zarr(path):
     try:
         store = zarr.ZipStore(path, mode="r")
     except zipfile.BadZipFile as bzf:
-        raise exceptions.FileFormatError("File is not in tgzip format") from bzf
+        raise exceptions.FileFormatError("File is not in tszip format") from bzf
     root = zarr.group(store=store)
     try:
         check_format(root)