Merge pull request #74 from swen128/fix-csv-processor

let CsvFileProcessor.dump accepts a pd.Series argument

Author: nishiba

gokart/file_processor.py

@@ -104,8 +104,9 @@ def load(self, file):
             return pd.DataFrame()
 
     def dump(self, obj, file):
-        assert isinstance(obj, pd.DataFrame), f'requires pd.DataFrame, but {type(obj)} is passed.'
-        obj.to_csv(file, index=False, sep=self._sep)
+        assert isinstance(obj, (pd.DataFrame, pd.Series)), \
+            f'requires pd.DataFrame or pd.Series, but {type(obj)} is passed.'
+        obj.to_csv(file, index=False, sep=self._sep, header=True)
 
 
 class GzipFileProcessor(FileProcessor):

test/test_target.py

@@ -83,6 +83,16 @@ def test_last_modified_time_without_file(self):
         with self.assertRaises(FileNotFoundError):
             target.last_modification_time()
 
+    def test_save_pandas_series(self):
+        obj = pd.Series(data=[1, 2], name='column_name')
+        file_path = os.path.join(_get_temporary_directory(), 'test.csv')
+
+        target = make_target(file_path=file_path, unique_id=None)
+        target.dump(obj)
+        loaded = target.load()
+
+        pd.testing.assert_series_equal(loaded['column_name'], obj)
+
 
 class S3TargetTest(unittest.TestCase):
     @mock_s3