make it possible to pass file processor

m3dev · Jul 13, 2019 · 122380d · 122380d
1 parent 63e1dab
commit 122380d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/gokart/target.py b/gokart/target.py
@@ -138,8 +138,8 @@ def load(file_path: str) -> pd.DataFrame:
         return pd.concat([pd.read_pickle(file_path) for file_path in glob(os.path.join(dir_path, 'data_*.pkl'))])
 
 
-def _make_file_system_target(file_path: str) -> luigi.target.FileSystemTarget:
-    processor = make_file_processor(file_path)
+def _make_file_system_target(file_path: str, processor: Optional[FileProcessor] = None) -> luigi.target.FileSystemTarget:
+    processor = processor or make_file_processor(file_path)
     if file_path.startswith('s3://'):
         return luigi.contrib.s3.S3Target(file_path, client=S3Config().get_s3_client(), format=processor.format())
     return luigi.LocalTarget(file_path, format=processor.format())
@@ -163,7 +163,7 @@ def _get_last_modification_time(path: str) -> datetime:
 def make_target(file_path: str, unique_id: Optional[str] = None, processor: Optional[FileProcessor] = None) -> TargetOnKart:
     file_path = _make_file_path(file_path, unique_id)
     processor = processor or make_file_processor(file_path)
-    file_system_target = _make_file_system_target(file_path)
+    file_system_target = _make_file_system_target(file_path, processor=processor)
     return SingleFileTarget(target=file_system_target, processor=processor)
 
 

diff --git a/test/test_task_on_kart.py b/test/test_task_on_kart.py
@@ -115,7 +115,7 @@ def test_make_target_without_id(self):
     def test_make_target_with_processor(self):
         task = _DummyTask()
         processor = XmlFileProcessor()
-        target = task.make_target('test.txt', processor=processor)
+        target = task.make_target('test.dummy', processor=processor)
         self.assertEqual(target._processor, processor)
         self.assertIsInstance(target, SingleFileTarget)