diff options
| -rw-r--r-- | gemato/compression.py | 20 | ||||
| -rw-r--r-- | tests/test_compression.py | 17 |
2 files changed, 34 insertions, 3 deletions
diff --git a/gemato/compression.py b/gemato/compression.py index 51850f7..4a3c9c6 100644 --- a/gemato/compression.py +++ b/gemato/compression.py @@ -99,8 +99,8 @@ def open_potentially_compressed_path(path, mode, **kwargs): Returns an object that must be used via the context manager API. """ - base, ext = os.path.splitext(path) - if ext not in ('.gz', '.bz2', '.lzma', '.xz'): + compression = get_compressed_suffix_from_filename(path) + if compression is None: return io.open(path, mode, **kwargs) bmode = mode @@ -110,7 +110,8 @@ def open_potentially_compressed_path(path, mode, **kwargs): f = io.open(path, bmode) fs = FileStack([f]) try: - cf = open_compressed_file(ext[1:], f, bmode if kwargs else mode) + cf = open_compressed_file(compression, f, + bmode if kwargs else mode) fs.files.append(cf) # add a TextIOWrapper on top whenever we do not want @@ -132,3 +133,16 @@ def get_potential_compressed_names(path): """ return [path + x for x in ('', '.gz', '.bz2', '.lzma', '.xz')] + + +def get_compressed_suffix_from_filename(path): + """ + Get the appropriate suffix (suitable for open_compressed_file()) + for a potentially compressed filename @path. If the path + does not seem to be compressed, returns None. + """ + + base, ext = os.path.splitext(path) + if ext in ('.gz', '.bz2', '.lzma', '.xz'): + return ext[1:] + return None diff --git a/tests/test_compression.py b/tests/test_compression.py index 5f614d3..df78ec8 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -658,3 +658,20 @@ class OtherUtilityTests(unittest.TestCase): 'test.lzma', 'test.xz', ])) + + def test_get_compressed_suffix_from_filename(self): + self.assertEqual( + gemato.compression.get_compressed_suffix_from_filename( + 'test.gz'), 'gz') + self.assertEqual( + gemato.compression.get_compressed_suffix_from_filename( + 'test.bz2'), 'bz2') + self.assertEqual( + gemato.compression.get_compressed_suffix_from_filename( + 'test.lzma'), 'lzma') + self.assertEqual( + gemato.compression.get_compressed_suffix_from_filename( + 'test.xz'), 'xz') + self.assertIsNone( + gemato.compression.get_compressed_suffix_from_filename( + 'test')) |