feat: add TaskOnKart.dump type

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):
@@ -42,13 +44,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)}')
@@ -66,15 +68,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

@@ -4,7 +4,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
@@ -24,7 +24,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.
 
@@ -269,7 +272,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:
@@ -305,7 +308,15 @@ def _flatten_recursively(dfs):
             data = data[list(required_columns)]
         return data
 
-    def dump(self, obj, target: Union[None, str, TargetOnKart] = None) -> None:
+    @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):
             assert not obj.empty
@@ -323,13 +334,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]
 # All the rules are listed on https://docs.astral.sh/ruff/rules/