pass expected_output_dataframe_type to make_target

Author: Ping-Han Hsieh

gokart/target.py

@@ -227,12 +227,13 @@ def make_target(
     processor: Optional[FileProcessor] = None,
     task_lock_params: Optional[TaskLockParams] = None,
     store_index_in_feather: bool = True,
+    expected_dataframe_type: Optional[pa.DataFrameModel] = None,
 ) -> 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, expected_dataframe_type=expected_dataframe_type)
 
 
 def make_model_target(

gokart/task.py

@@ -8,6 +8,7 @@
 
 import luigi
 import pandas as pd
+import pandera as pa
 from luigi.parameter import ParameterVisibility
 
 import gokart
@@ -83,6 +84,7 @@ class TaskOnKart(luigi.Task):
         default=True, description='Check if output file exists at run. If exists, run() will be skipped.', significant=False
     )
     should_lock_run: bool = ExplicitBoolParameter(default=False, significant=False, description='Whether to use redis lock or not at task run.')
+    expected_output_dataframe_type: Optional[pa.DataFrameModel] = None
 
     def __init__(self, *args, **kwargs):
         self._add_configuration(kwargs, 'TaskOnKart')
@@ -192,7 +194,7 @@ 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, expected_dataframe_type=self.expected_output_dataframe_type
         )
 
     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_s3
 

test/test_task_on_kart.py

@@ -7,6 +7,7 @@
 
 import luigi
 import pandas as pd
+import pandera as pa
 from luigi.parameter import ParameterVisibility
 from luigi.util import inherits
 
@@ -340,7 +341,19 @@ def test_fail_on_empty_dump(self):
         # fail
         task = _DummyTask(fail_on_empty_dump=True)
         self.assertRaises(AssertionError, lambda: task.dump(pd.DataFrame()))
-
+    
+    def test_fail_with_type_check(self):
+        
+        class _DummyTypeSchema(pa.DataFrameModel):
+            a: pa.typing.Series[int] = pa.Field()
+        class _DummyTaskWithType(gokart.TaskOnKart):
+            expected_output_dataframe_type = _DummyTypeSchema
+
+        task = _DummyTaskWithType()
+        # fail
+        with self.assertRaises(pa.errors.SchemaError):
+            task.dump(pd.DataFrame(dict(a=['1', '2', '3'])))
+    
     @patch('luigi.configuration.get_config')
     def test_add_configuration(self, mock_config: MagicMock):
         mock_config.return_value = {'_DummyTask': {'list_param': '["c", "d"]', 'param': '3', 'bool_param': 'True'}}