feat: add TaskOnKart.dump type (#368)

gokart/build.py

@@ -1,7 +1,7 @@
 import logging
 from functools import partial
 from logging import getLogger
-from typing import Any, Optional
+from typing import Literal, Optional, TypeVar, cast, overload
 
 import backoff
 import luigi
@@ -11,6 +11,8 @@
 from gokart.target import TargetOnKart
 from gokart.task import TaskOnKart
 
+T = TypeVar('T')
+
 
 class LoggerConfig:
     def __init__(self, level: int):
@@ -41,13 +43,13 @@ def __init__(self):
         self.flag: bool = False
 
 
-def _get_output(task: TaskOnKart) -> Any:
+def _get_output(task: TaskOnKart[T]) -> T:
     output = task.output()
     # FIXME: currently, nested output is not supported
     if isinstance(output, list) or isinstance(output, tuple):
-        return [t.load() for t in output if isinstance(t, TargetOnKart)]
+        return cast(T, [t.load() for t in output if isinstance(t, TargetOnKart)])
     if isinstance(output, dict):
-        return {k: t.load() for k, t in output.items() if isinstance(t, TargetOnKart)}
+        return cast(T, {k: t.load() for k, t in output.items() if isinstance(t, TargetOnKart)})
     if isinstance(output, TargetOnKart):
         return output.load()
     raise ValueError(f'output type is not supported: {type(output)}')
@@ -65,15 +67,39 @@ def _reset_register(keep={'gokart', 'luigi'}):
     ]
 
 
+@overload
+def build(
+    task: TaskOnKart[T],
+    return_value: Literal[True] = True,
+    reset_register: bool = True,
+    log_level: int = logging.ERROR,
+    task_lock_exception_max_tries: int = 10,
+    task_lock_exception_max_wait_seconds: int = 600,
+    **env_params,
+) -> T: ...
+
+
+@overload
+def build(
+    task: TaskOnKart[T],
+    return_value: Literal[False],
+    reset_register: bool = True,
+    log_level: int = logging.ERROR,
+    task_lock_exception_max_tries: int = 10,
+    task_lock_exception_max_wait_seconds: int = 600,
+    **env_params,
+) -> None: ...
+
+
 def build(
-    task: TaskOnKart,
+    task: TaskOnKart[T],
     return_value: bool = True,
     reset_register: bool = True,
     log_level: int = logging.ERROR,
     task_lock_exception_max_tries: int = 10,
     task_lock_exception_max_wait_seconds: int = 600,
     **env_params,
-) -> Optional[Any]:
+) -> Optional[T]:
     """
     Run gokart task for local interpreter.
     Sharing the most of its parameters with luigi.build (see https://luigi.readthedocs.io/en/stable/api/luigi.html?highlight=build#luigi.build)

gokart/task.py

@@ -5,7 +5,7 @@
 import types
 from importlib import import_module
 from logging import getLogger
-from typing import Any, Callable, Dict, List, Optional, Set, Union
+from typing import Any, Callable, Dict, Generator, Generic, List, Optional, Set, TypeVar, Union, overload
 
 import luigi
 import pandas as pd
@@ -25,7 +25,10 @@
 logger = getLogger(__name__)
 
 
-class TaskOnKart(luigi.Task):
+T = TypeVar('T')
+
+
+class TaskOnKart(luigi.Task, Generic[T]):
     """
     This is a wrapper class of luigi.Task.
 
@@ -282,7 +285,7 @@ def _load(targets):
             return list(data.values())[0]
         return data
 
-    def load_generator(self, target: Union[None, str, TargetOnKart] = None) -> Any:
+    def load_generator(self, target: Union[None, str, TargetOnKart] = None) -> Generator[Any, None, None]:
         def _load(targets):
             if isinstance(targets, list) or isinstance(targets, tuple):
                 for t in targets:
@@ -318,6 +321,12 @@ def _flatten_recursively(dfs):
             data = data[list(required_columns)]
         return data
 
+    @overload
+    def dump(self, obj: T, target: None = None) -> None: ...
+
+    @overload
+    def dump(self, obj: Any, target: Union[str, TargetOnKart]) -> None: ...
+
     def dump(self, obj: Any, target: Union[None, str, TargetOnKart] = None) -> None:
         PandasTypeConfigMap().check(obj, task_namespace=self.task_namespace)
         if self.fail_on_empty_dump and isinstance(obj, pd.DataFrame):
@@ -336,13 +345,13 @@ def get_own_code(self):
         own_codes = self.get_code(self)
         return ''.join(sorted(list(own_codes - gokart_codes)))
 
-    def make_unique_id(self):
+    def make_unique_id(self) -> str:
         unique_id = self.task_unique_id or self._make_hash_id()
         if self.cache_unique_id:
             self.task_unique_id = unique_id
         return unique_id
 
-    def _make_hash_id(self):
+    def _make_hash_id(self) -> str:
         def _to_str_params(task):
             if isinstance(task, TaskOnKart):
                 return str(task.make_unique_id()) if task.significant else None

pyproject.toml

@@ -43,6 +43,7 @@ fakeredis = "*"
 mypy = "*"
 types-redis = "*"
 matplotlib = "*"
+typing-extensions = "^4.11.0"
 
 [tool.ruff]
 line-length = 160

test/test_build.py

@@ -1,7 +1,14 @@
 import logging
 import os
+import sys
 import unittest
 from copy import copy
+from typing import Dict
+
+if sys.version_info >= (3, 11):
+    from typing import assert_type
+else:
+    from typing_extensions import assert_type
 
 import luigi
 import luigi.mock
@@ -11,9 +18,9 @@
 from gokart.conflict_prevention_lock.task_lock import TaskLockException
 
 
-class _DummyTask(gokart.TaskOnKart):
+class _DummyTask(gokart.TaskOnKart[str]):
     task_namespace = __name__
-    param = luigi.Parameter()
+    param: str = luigi.Parameter()
 
     def output(self):
         return self.make_target('./test/dummy.pkl')
@@ -22,10 +29,10 @@ def run(self):
         self.dump(self.param)
 
 
-class _DummyTaskTwoOutputs(gokart.TaskOnKart):
+class _DummyTaskTwoOutputs(gokart.TaskOnKart[Dict[str, str]]):
     task_namespace = __name__
-    param1 = luigi.Parameter()
-    param2 = luigi.Parameter()
+    param1: str = luigi.Parameter()
+    param2: str = luigi.Parameter()
 
     def output(self):
         return {'out1': self.make_target('./test/dummy1.pkl'), 'out2': self.make_target('./test/dummy2.pkl')}
@@ -42,7 +49,7 @@ def run(self):
         raise RuntimeError
 
 
-class _ParallelRunner(gokart.TaskOnKart):
+class _ParallelRunner(gokart.TaskOnKart[str]):
     def requires(self):
         return [_DummyTask(param=str(i)) for i in range(10)]
 
@@ -79,6 +86,7 @@ def test_read_config(self):
         config_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config', 'test_config.ini')
         gokart.utils.add_config(config_file_path)
         output = gokart.build(_DummyTask(), reset_register=False)
+        assert_type(output, str)
         self.assertEqual(output, 'test')
 
     def test_build_dict_outputs(self):
@@ -87,6 +95,7 @@ def test_build_dict_outputs(self):
             'out2': 'test2',
         }
         output = gokart.build(_DummyTaskTwoOutputs(param1=param_dict['out1'], param2=param_dict['out2']), reset_register=False)
+        assert_type(output, Dict[str, str])
         self.assertEqual(output, param_dict)
 
     def test_failed_task(self):