Skip to content

Commit 9149f85

Browse files
author
Ping-Han Hsieh
committed
fix ruff
1 parent f1dac48 commit 9149f85

File tree

4 files changed

+17
-10
lines changed

4 files changed

+17
-10
lines changed

gokart/target.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,7 @@ def make_target(
227227
processor: Optional[FileProcessor] = None,
228228
task_lock_params: Optional[TaskLockParams] = None,
229229
store_index_in_feather: bool = True,
230-
expected_dataframe_type: Optional[pa.DataFrameModel] = None,
230+
expected_dataframe_type: Optional[type[pa.DataFrameModel]] = None,
231231
) -> TargetOnKart:
232232
_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)
233233
file_path = _make_file_path(file_path, unique_id)

gokart/task.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ class TaskOnKart(luigi.Task):
8484
default=True, description='Check if output file exists at run. If exists, run() will be skipped.', significant=False
8585
)
8686
should_lock_run: bool = ExplicitBoolParameter(default=False, significant=False, description='Whether to use redis lock or not at task run.')
87-
expected_output_dataframe_type: Optional[pa.DataFrameModel] = None
87+
expected_output_dataframe_type: Optional[type[pa.DataFrameModel]] = None
8888

8989
def __init__(self, *args, **kwargs):
9090
self._add_configuration(kwargs, 'TaskOnKart')
@@ -194,7 +194,12 @@ def make_target(self, relative_file_path: Optional[str] = None, use_unique_id: b
194194
)
195195

196196
return gokart.target.make_target(
197-
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
197+
file_path=file_path,
198+
unique_id=unique_id,
199+
processor=processor,
200+
task_lock_params=task_lock_params,
201+
store_index_in_feather=self.store_index_in_feather,
202+
expected_dataframe_type=self.expected_output_dataframe_type,
198203
)
199204

200205
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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -269,14 +269,15 @@ def test_typed_target(self):
269269
_task_lock_params = make_task_lock_params(file_path=file_path, unique_id=unique_id)
270270
processor = make_file_processor(file_path, store_index_in_feather=False)
271271
file_system_target = luigi.LocalTarget(file_path, format=processor.format())
272-
file_target = SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema)
272+
file_target = SingleFileTarget(
273+
target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema
274+
)
273275

274276
file_target.dump(test_case)
275277
dumped_data = file_target.load()
276278
self.assertIsInstance(dumped_data, pa.typing.DataFrame)
277279
self.DummyDataFrameSchema.validate(dumped_data)
278280

279-
280281
def test_invalid_typed_target(self):
281282
test_case = pd.DataFrame(dict(a=['1', '2']))
282283

@@ -286,11 +287,12 @@ def test_invalid_typed_target(self):
286287
_task_lock_params = make_task_lock_params(file_path=file_path, unique_id=unique_id)
287288
processor = make_file_processor(file_path, store_index_in_feather=False)
288289
file_system_target = luigi.LocalTarget(file_path, format=processor.format())
289-
file_target = SingleFileTarget(target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema)
290+
file_target = SingleFileTarget(
291+
target=file_system_target, processor=processor, task_lock_params=_task_lock_params, expected_dataframe_type=self.DummyDataFrameSchema
292+
)
290293

291294
with self.assertRaises(pa.errors.SchemaError):
292295
file_target.dump(test_case)
293-
294296

295297

296298
if __name__ == '__main__':

test/test_task_on_kart.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -341,19 +341,19 @@ def test_fail_on_empty_dump(self):
341341
# fail
342342
task = _DummyTask(fail_on_empty_dump=True)
343343
self.assertRaises(AssertionError, lambda: task.dump(pd.DataFrame()))
344-
344+
345345
def test_fail_with_type_check(self):
346-
347346
class _DummyTypeSchema(pa.DataFrameModel):
348347
a: pa.typing.Series[int] = pa.Field()
348+
349349
class _DummyTaskWithType(gokart.TaskOnKart):
350350
expected_output_dataframe_type = _DummyTypeSchema
351351

352352
task = _DummyTaskWithType()
353353
# fail
354354
with self.assertRaises(pa.errors.SchemaError):
355355
task.dump(pd.DataFrame(dict(a=['1', '2', '3'])))
356-
356+
357357
@patch('luigi.configuration.get_config')
358358
def test_add_configuration(self, mock_config: MagicMock):
359359
mock_config.return_value = {'_DummyTask': {'list_param': '["c", "d"]', 'param': '3', 'bool_param': 'True'}}

0 commit comments

Comments
 (0)