Merge pull request #1250: Enforce types to io.open_file

CHANGES.md

@@ -9,9 +9,11 @@
 ### Bug fixes
 
 * parse: Fix a bug where `--fix-dates` was always applied, with a default of `--fix-dates=monthfirst`. Now, running without `--fix-dates` will leave dates as-is. [#1247][] (@victorlin)
+* `augur.io.open_file`: Previously, the docs described a type restriction on `path_or_buffer` but it was not enforced. It has been updated to allow all I/O classes, and is enforced at run-time. [#1250][] (@victorlin)
 
 [#1240]: https://github.com/nextstrain/augur/pull/1240
 [#1247]: https://github.com/nextstrain/augur/issues/1247
+[#1250]: https://github.com/nextstrain/augur/pull/1250
 
 ## 22.0.3 (14 June 2023)
 

augur/io/file.py

@@ -1,5 +1,7 @@
+import os
 from contextlib import contextmanager
-from xopen import xopen
+from io import IOBase
+from xopen import PipedCompressionReader, PipedCompressionWriter, xopen
 
 
 @contextmanager
@@ -10,7 +12,7 @@ def open_file(path_or_buffer, mode="r", **kwargs):
 
     Parameters
     ----------
-    path_or_buffer : str or `os.PathLike` or `io.StringIO`
+    path_or_buffer
         Name of the file to open or an existing IO buffer
 
     mode : str
@@ -22,8 +24,12 @@ def open_file(path_or_buffer, mode="r", **kwargs):
         File handle object
 
     """
-    try:
+    if isinstance(path_or_buffer, (str, os.PathLike)):
         with xopen(path_or_buffer, mode, **kwargs) as handle:
             yield handle
-    except TypeError:
+
+    elif isinstance(path_or_buffer, (IOBase, PipedCompressionReader, PipedCompressionWriter)):
         yield path_or_buffer
+
+    else:
+        raise TypeError(f"Type {type(path_or_buffer)} is not supported.")

tests/io/test_file.py

@@ -1,3 +1,4 @@
+import pytest
 import augur.io.file
 
 
@@ -71,3 +72,46 @@ def test_open_file_write_zstd(self, tmpdir):
             f_write.write('foo\nbar\n')
         with zstd.open(path, 'rt') as f_read:
             assert f_read.read() == 'foo\nbar\n'
+
+    def test_open_file_nested_read(self, tmpdir):
+        """Open a text file as a path then buffer for reading."""
+        path = str(tmpdir / 'test.txt')
+        with open(path, 'wt') as f_write:
+            f_write.write('foo\nbar\n')
+
+        with augur.io.file.open_file(path) as f_read:
+            with augur.io.file.open_file(f_read) as f_read:
+                assert f_read.read() == 'foo\nbar\n'
+
+    def test_open_file_nested_read_zstd(self, tmpdir):
+        """Open a zstd-compressed text file as a path then buffer for reading."""
+        import zstandard as zstd
+        path = str(tmpdir / 'test.txt.zst')
+        with zstd.open(path, 'wt') as f_write:
+            f_write.write('foo\nbar\n')
+
+        with augur.io.file.open_file(path) as f_read:
+            with augur.io.file.open_file(f_read) as f_read:
+                assert f_read.read() == 'foo\nbar\n'
+
+    def test_open_file_nested_write(self, tmpdir):
+        """Open a text file as a path then buffer for writing."""
+        path = str(tmpdir / 'test.txt')
+        with augur.io.file.open_file(path, 'w') as f_write:
+            with augur.io.file.open_file(f_write, 'w') as f_write:
+                f_write.write('foo\nbar\n')
+
+    def test_open_file_nested_write_zstd(self, tmpdir):
+        """Open a zstd-compressed text file as a path then buffer for writing."""
+        path = str(tmpdir / 'test.txt.zst')
+        with augur.io.file.open_file(path, 'w') as f_write:
+            with augur.io.file.open_file(f_write, 'w') as f_write:
+                f_write.write('foo\nbar\n')
+
+    def test_open_file_read_error(self):
+        """Try reading from an unsupported type."""
+        path_or_buffer = len("bogus")
+
+        with pytest.raises(TypeError, match="Type <class 'int'> is not supported."):
+            with augur.io.file.open_file(path_or_buffer, 'r') as f:
+                f.read()