check_round_trip refactoring to trap on FastParquet wrires to s3.

Author: maxim veksler

pandas/tests/io/test_parquet.py

@@ -212,28 +212,37 @@ def check_error_on_write(self, df, engine, exc):
             with tm.ensure_clean() as path:
                 to_parquet(df, path, engine, compression=None)
 
-    def check_round_trip(self, df, engine, expected=None,
-                         write_kwargs=None, read_kwargs=None,
-                         check_names=True):
+    def do_round_trip(self, df, path, engine_impl, expected=None,
+                      write_kwargs=None, read_kwargs=None,
+                      check_names=True):
+
         if write_kwargs is None:
-            write_kwargs = {}
+            write_kwargs = {'compression': None}
+
         if read_kwargs is None:
             read_kwargs = {}
-        with tm.ensure_clean() as path:
-            df.to_parquet(path, engine, **write_kwargs)
-            result = read_parquet(path, engine, **read_kwargs)
 
-            if expected is None:
-                expected = df
-            tm.assert_frame_equal(result, expected, check_names=check_names)
+        df.to_parquet(path, engine_impl, **write_kwargs)
+        actual = read_parquet(path, engine_impl, **read_kwargs)
 
-            # repeat
-            to_parquet(df, path, engine, **write_kwargs)
-            result = pd.read_parquet(path, engine, **read_kwargs)
+        if expected is None:
+            expected = df
 
-            if expected is None:
-                expected = df
-            tm.assert_frame_equal(result, expected, check_names=check_names)
+        tm.assert_frame_equal(expected, actual, check_names=check_names)
+
+    def check_round_trip(self, df, engine, expected=None,
+                         write_kwargs=None, read_kwargs=None,
+                         check_names=True):
+
+        with tm.ensure_clean() as path:
+            self.do_round_trip(df, path, engine, expected,
+                               write_kwargs=write_kwargs, read_kwargs=read_kwargs,
+                               check_names=check_names)
+
+            # repeat
+            self.do_round_trip(df, path, engine, expected,
+                               write_kwargs=write_kwargs, read_kwargs=read_kwargs,
+                               check_names=check_names)
 
 
 class TestBasic(Base):
@@ -251,7 +260,7 @@ def test_columns_dtypes(self, engine):
 
         # unicode
         df.columns = [u'foo', u'bar']
-        self.check_round_trip(df, engine, write_kwargs={'compression': None})
+        self.check_round_trip(df, engine)
 
     def test_columns_dtypes_invalid(self, engine):
 
@@ -292,7 +301,6 @@ def test_read_columns(self, engine):
 
         expected = pd.DataFrame({'string': list('abc')})
         self.check_round_trip(df, engine, expected=expected,
-                              write_kwargs={'compression': None},
                               read_kwargs={'columns': ['string']})
 
     def test_write_index(self, engine):
@@ -304,7 +312,7 @@ def test_write_index(self, engine):
                 pytest.skip("pyarrow is < 0.7.0")
 
         df = pd.DataFrame({'A': [1, 2, 3]})
-        self.check_round_trip(df, engine, write_kwargs={'compression': None})
+        self.check_round_trip(df, engine)
 
         indexes = [
             [2, 3, 4],
@@ -315,15 +323,12 @@ def test_write_index(self, engine):
         # non-default index
         for index in indexes:
             df.index = index
-            self.check_round_trip(
-                df, engine,
-                write_kwargs={'compression': None},
-                check_names=check_names)
+            self.check_round_trip(df, engine, check_names=check_names)
 
         # index with meta-data
         df.index = [0, 1, 2]
         df.index.name = 'foo'
-        self.check_round_trip(df, engine, write_kwargs={'compression': None})
+        self.check_round_trip(df, engine)
 
     def test_write_multiindex(self, pa_ge_070):
         # Not suppoprted in fastparquet as of 0.1.3 or older pyarrow version
@@ -332,7 +337,7 @@ def test_write_multiindex(self, pa_ge_070):
         df = pd.DataFrame({'A': [1, 2, 3]})
         index = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)])
         df.index = index
-        self.check_round_trip(df, engine, write_kwargs={'compression': None})
+        self.check_round_trip(df, engine)
 
     def test_write_column_multiindex(self, engine):
         # column multi-index
@@ -428,13 +433,7 @@ def test_categorical_unsupported(self, pa_lt_070):
 
     def test_s3_roundtrip(self, df_compat, s3_resource, pa):
         # GH #19134
-        df_compat.to_parquet('s3://pandas-test/test.parquet',
-                             engine=pa, compression=None)
-
-        expected = df_compat
-        actual = read_parquet('s3://pandas-test/test.parquet', engine=pa)
-
-        tm.assert_frame_equal(expected, actual)
+        self.do_round_trip(df_compat, 's3://pandas-test/test.parquet', pa)
 
 
 class TestParquetFastParquet(Base):
@@ -446,7 +445,7 @@ def test_basic(self, fp, df_full):
         # additional supported types for fastparquet
         df['timedelta'] = pd.timedelta_range('1 day', periods=3)
 
-        self.check_round_trip(df, fp, write_kwargs={'compression': None})
+        self.check_round_trip(df, fp)
 
     @pytest.mark.skip(reason="not supported")
     def test_duplicate_columns(self, fp):
@@ -459,8 +458,7 @@ def test_duplicate_columns(self, fp):
     def test_bool_with_none(self, fp):
         df = pd.DataFrame({'a': [True, None, False]})
         expected = pd.DataFrame({'a': [1.0, np.nan, 0.0]}, dtype='float16')
-        self.check_round_trip(df, fp, expected=expected,
-                              write_kwargs={'compression': None})
+        self.check_round_trip(df, fp, expected=expected)
 
     def test_unsupported(self, fp):
 
@@ -476,7 +474,7 @@ def test_categorical(self, fp):
         if LooseVersion(fastparquet.__version__) < LooseVersion("0.1.3"):
             pytest.skip("CategoricalDtype not supported for older fp")
         df = pd.DataFrame({'a': pd.Categorical(list('abc'))})
-        self.check_round_trip(df, fp, write_kwargs={'compression': None})
+        self.check_round_trip(df, fp)
 
     def test_datetime_tz(self, fp):
         # doesn't preserve tz
@@ -485,8 +483,7 @@ def test_datetime_tz(self, fp):
 
         # warns on the coercion
         with catch_warnings(record=True):
-            self.check_round_trip(df, fp, df.astype('datetime64[ns]'),
-                                  write_kwargs={'compression': None})
+            self.check_round_trip(df, fp, df.astype('datetime64[ns]'))
 
     def test_filter_row_groups(self, fp):
         d = {'a': list(range(0, 3))}
@@ -497,3 +494,10 @@ def test_filter_row_groups(self, fp):
             result = read_parquet(path, fp, filters=[('a', '==', 0)])
         assert len(result) == 1
 
+    def test_s3_roundtrip(self, df_compat, s3_resource, fp):
+        print(s3_resource, fp)
+
+        # GH #19134
+        with pytest.raises(TypeError):
+            self.do_round_trip(df_compat, 's3://pandas-test/test.parquet', fp)
+