Test _infer_compression in io/test_common.py

Author: dhimmel

pandas/tests/io/parser/compression.py

@@ -157,18 +157,6 @@ def test_read_csv_infer_compression(self):
 
         inputs[3].close()
 
-    @pytest.mark.parametrize('extension', ['', '.gz', '.bz2'])
-    def test_read_csv_infer_compression_pathlib(self, extension):
-        """
-        Test that compression is inferred from pathlib.Path paths.
-        """
-        pathlib = pytest.importorskip('pathlib')
-        read_csv_kwargs = {'index_col': 0, 'parse_dates': True}
-        expected = self.read_csv(self.csv1, **read_csv_kwargs)
-        path = pathlib.Path(self.csv1 + extension)
-        df = self.read_csv(path, compression='infer', **read_csv_kwargs)
-        tm.assert_frame_equal(expected, df)
-
     def test_invalid_compression(self):
         msg = 'Unrecognized compression type: sfark'
         with tm.assert_raises_regex(ValueError, msg):

pandas/tests/io/test_common.py

@@ -14,16 +14,6 @@
 
 from pandas import read_csv, concat
 
-try:
-    from pathlib import Path
-except ImportError:
-    pass
-
-try:
-    from py.path import local as LocalPath
-except ImportError:
-    pass
-
 
 class CustomFSPath(object):
     """For testing fspath on unknown objects"""
@@ -34,6 +24,21 @@ def __fspath__(self):
         return self.path
 
 
+# Functions that consume a string path and return a string or path-like object
+path_types = [str, CustomFSPath]
+
+try:
+    from pathlib import Path
+    path_types.append(Path)
+except ImportError:
+    pass
+
+try:
+    from py.path import local as LocalPath
+    path_types.append(LocalPath)
+except ImportError:
+    pass
+
 HERE = os.path.dirname(__file__)
 
 
@@ -83,6 +88,19 @@ def test_stringify_path_fspath(self):
         result = common._stringify_path(p)
         assert result == 'foo/bar.csv'
 
+    @pytest.mark.parametrize('extension,expected', [
+        ('', None),
+        ('.gz', 'gzip'),
+        ('.bz2', 'bz2'),
+        ('.zip', 'zip'),
+        ('.xz', 'xz'),
+    ])
+    @pytest.mark.parametrize('path_type', path_types)
+    def test_infer_compression_from_path(self, extension, expected, path_type):
+        path = path_type('foo/bar.csv' + extension)
+        compression = common._infer_compression(path, compression='infer')
+        assert compression == expected
+
     def test_get_filepath_or_buffer_with_path(self):
         filename = '~/sometest'
         filepath_or_buffer, _, _ = common.get_filepath_or_buffer(filename)