feat: add validator for dumped object type check

Author: Shogo Ujiie

gokart/target.py

@@ -5,12 +5,11 @@
 from datetime import datetime
 from glob import glob
 from logging import getLogger
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import luigi
 import numpy as np
 import pandas as pd
-import pandera as pa
 
 from gokart.conflict_prevention_lock.task_lock import TaskLockParams, make_task_lock_params
 from gokart.conflict_prevention_lock.task_lock_wrappers import wrap_dump_with_lock, wrap_load_with_lock, wrap_remove_with_lock
@@ -79,11 +78,12 @@ def __init__(
         target: luigi.target.FileSystemTarget,
         processor: FileProcessor,
         task_lock_params: TaskLockParams,
-        expected_dataframe_type: Optional[pa.DataFrameModel] = None,
+        validator: Callable[[Any], bool] = lambda x: True,
     ) -> None:
         self._target = target
         self._processor = processor
         self._task_lock_params = task_lock_params
+        self._validator = validator
 
     def _exists(self) -> bool:
         return self._target.exists()
@@ -94,14 +94,12 @@ def _get_task_lock_params(self) -> TaskLockParams:
     def _load(self) -> Any:
         with self._target.open('r') as f:
             obj = self._processor.load(f)
-            if self.expected_dataframe_type is not None:
-                return self.expected_dataframe_type(obj)
+            self._validator(obj)
 
             return obj
 
     def _dump(self, obj) -> None:
-        if self.expected_dataframe_type is not None:
-            self.expected_dataframe_type.validate(obj)
+        self._validator(obj)
 
         with self._target.open('w') as f:
             self._processor.dump(obj, f)
@@ -225,12 +223,13 @@ def make_target(
     processor: Optional[FileProcessor] = None,
     task_lock_params: Optional[TaskLockParams] = None,
     store_index_in_feather: bool = True,
+    validator: Callable[[Any], bool] = lambda x: True,
 ) -> TargetOnKart:
     _task_lock_params = task_lock_params if task_lock_params is not None else make_task_lock_params(file_path=file_path, unique_id=unique_id)
     file_path = _make_file_path(file_path, unique_id)
     processor = processor or make_file_processor(file_path, store_index_in_feather=store_index_in_feather)
     file_system_target = _make_file_system_target(file_path, processor=processor, store_index_in_feather=store_index_in_feather)
-    return SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params)
+    return SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, validator=validator)
 
 
 def make_model_target(

gokart/task.py

@@ -193,7 +193,13 @@ def clone(self, cls=None, **kwargs):
 
         return cls(**new_k)
 
-    def make_target(self, relative_file_path: Optional[str] = None, use_unique_id: bool = True, processor: Optional[FileProcessor] = None) -> TargetOnKart:
+    def make_target(
+        self,
+        relative_file_path: Optional[str] = None,
+        use_unique_id: bool = True,
+        processor: Optional[FileProcessor] = None,
+        validator: Callable[[Any], bool] = lambda x: True,
+    ) -> TargetOnKart:
         formatted_relative_file_path = (
             relative_file_path if relative_file_path is not None else os.path.join(self.__module__.replace('.', '/'), f'{type(self).__name__}.pkl')
         )
@@ -210,7 +216,12 @@ def make_target(self, relative_file_path: Optional[str] = None, use_unique_id: b
         )
 
         return gokart.target.make_target(
-            file_path=file_path, unique_id=unique_id, processor=processor, task_lock_params=task_lock_params, store_index_in_feather=self.store_index_in_feather
+            file_path=file_path,
+            unique_id=unique_id,
+            processor=processor,
+            task_lock_params=task_lock_params,
+            store_index_in_feather=self.store_index_in_feather,
+            validator=validator,
         )
 
     def make_large_data_frame_target(self, relative_file_path: Optional[str] = None, use_unique_id: bool = True, max_byte=int(2**26)) -> TargetOnKart:

test/test_target.py

@@ -9,8 +9,8 @@
 import boto3
 import luigi
 import numpy as np
-import pandera as pa
 import pandas as pd
+import pandera as pa
 from matplotlib import pyplot
 from moto import mock_aws
 
@@ -284,36 +284,40 @@ def test_model_target_on_s3(self):
 
 
 class SingleFileTargetTest(unittest.TestCase):
-    class DummyDataFrameSchema(pa.DataFrameModel):
-        a: pa.typing.Series[int] = pa.Field()
-
     def test_typed_target(self):
+        def validate_dataframe(x):
+            return isinstance(x, pd.DataFrame)
+
         test_case = pd.DataFrame(dict(a=[1, 2]))
 
         with tempfile.TemporaryDirectory() as temp_dir:
             _task_lock_params = None
-            file_path = os.path.join(temp_dir, 'test.csv')
+            file_path = os.path.join(temp_dir, 'test.pkl')
             processor = make_file_processor(file_path, store_index_in_feather=False)
             file_system_target = luigi.LocalTarget(file_path, format=processor.format())
-            file_target = SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema)
+            file_target = SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, validator=validate_dataframe)
 
             file_target.dump(test_case)
             dumped_data = file_target.load()
             self.assertIsInstance(dumped_data, self.DummyDataFrameSchema)
 
     def test_invalid_typed_target(self):
+        def validate_int(x):
+            return isinstance(x, int)
+
         test_case = pd.DataFrame(dict(a=['1', '2']))
 
         with tempfile.TemporaryDirectory() as temp_dir:
             _task_lock_params = None
             file_path = os.path.join(temp_dir, 'test.csv')
             processor = make_file_processor(file_path, store_index_in_feather=False)
             file_system_target = luigi.LocalTarget(file_path, format=processor.format())
-            file_target = SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema)
+            file_target = SingleFileTarget(
+                target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=validate_int
+            )
 
             with self.assertRaises(pa.errors.SchemaError):
                 file_target.dump(test_case)
-            
 
 
 if __name__ == '__main__':