feat: support type-safe Parameter annotations without mypy plugin (#496)

* feat: support generic type for TaskInstanceParameter and ListTaskInstanceParameter

* fix: unsupport non typed Task parameter

* fix: support new type hint

* docs: add next type hint

* chore: import from typing_extensions if Python version is less than 3.11

* fix: use StrParameter instead of Parameter

Author: kitagry

README.md

@@ -107,9 +107,9 @@ class TwoDumpTask(gokart.TaskOnKart[int]):
 
 class AddTask(gokart.TaskOnKart[int]):
     # `a` requires a task to dump `int`
-    a: gokart.TaskOnKart[int] = gokart.TaskInstanceParameter()
+    a: gokart.TaskInstanceParameter[gokart.TaskOnKart[int]] = gokart.TaskInstanceParameter()
     # `b` requires a task to dump `int`
-    b: gokart.TaskOnKart[int] = gokart.TaskInstanceParameter()
+    b: gokart.TaskInstanceParameter[gokart.TaskOnKart[int]] = gokart.TaskInstanceParameter()
 
     def requires(self):
         return dict(a=self.a, b=self.b)

docs/intro_to_gokart.rst

@@ -141,7 +141,7 @@ The :func:`~gokart.build` is inline code.
     import luigi
 
     class SampleTask(gokart.TaskOnKart[str]):
-        param = luigi.Parameter()
+        param: luigi.Parameter = luigi.Parameter()
 
         def run(self):
             self.dump(self.param)

docs/setting_task_parameters.rst

@@ -2,7 +2,7 @@
 Setting Task Parameters
 ============================
 
-There are several ways to set task parameters. 
+There are several ways to set task parameters.
 
 - Set parameter from command line
 - Set parameter at config file
@@ -26,7 +26,7 @@ Set parameter at config file
     [sample.SomeTask]
     param = Hello
 
-Above config file (``config.ini``) must be read before ``gokart.run()`` as the following code: 
+Above config file (``config.ini``) must be read before ``gokart.run()`` as the following code:
 
 .. code:: python
 
@@ -68,12 +68,12 @@ Parameter values can be inherited from other task using ``@inherits_config_param
 .. code:: python
 
     class MasterConfig(luigi.Config):
-        param: str = luigi.Parameter()
-        param2: str = luigi.Parameter()
+        param: luigi.Parameter = luigi.Parameter()
+        param2: luigi.Parameter = luigi.Parameter()
 
     @inherits_config_params(MasterConfig)
     class SomeTask(gokart.TaskOnKart):
-        param: str = luigi.Parameter()
+        param: luigi.Parameter = luigi.Parameter()
 
 
 This is useful when multiple tasks has the same parameter. In the above example, parameter settings of ``MasterConfig`` will be inherited to all tasks decorated with ``@inherits_config_params(MasterConfig)`` as ``SomeTask``.
@@ -84,12 +84,12 @@ In the above example, ``param2`` will not be available in ``SomeTask``, since ``
 .. code:: python
 
     class MasterConfig(luigi.Config):
-        param: str = luigi.Parameter()
-        param2: str = luigi.Parameter()
+        param: luigi.Parameter = luigi.Parameter()
+        param2: luigi.Parameter = luigi.Parameter()
 
     @inherits_config_params(MasterConfig, parameter_alias={'param2': 'param3'})
     class SomeTask(gokart.TaskOnKart):
-        param3: str = luigi.Parameter()
+        param3: luigi.Parameter = luigi.Parameter()
 
 
 You may also set a parameter name alias by setting ``parameter_alias``.

docs/task_on_kart.rst

@@ -17,7 +17,7 @@ How ``TaskOnKart`` helps to define a task looks like:
 
 
     class TaskA(gokart.TaskOnKart[str]):
-        param = luigi.Parameter()
+        param: luigi.Parameter = luigi.Parameter()
 
         def output(self):
             return self.make_target('output_of_task_a.pkl')
@@ -28,7 +28,7 @@ How ``TaskOnKart`` helps to define a task looks like:
 
 
     class TaskB(gokart.TaskOnKart[str]):
-        param = luigi.Parameter()
+        param: luigi.Parameter = luigi.Parameter()
 
         def requires(self):
             return TaskA(param='world')

docs/task_parameters.rst

@@ -11,9 +11,9 @@ Also please refer to :doc:`task_settings` section.
 .. code:: python
 
     class Task(gokart.TaskOnKart):
-        param_a = luigi.Parameter()
-        param_c = luigi.ListParameter()
-        param_d = luigi.IntParameter(default=1)
+        param_a: luigi.Parameter = luigi.Parameter()
+        param_c: luigi.ListParameter = luigi.ListParameter()
+        param_d: luigi.IntParameter = luigi.IntParameter(default=1)
 
 Please refer to `luigi document <https://luigi.readthedocs.io/en/stable/api/luigi.parameter.html>`_ for a list of parameter types.
 
@@ -42,7 +42,7 @@ The :func:`~gokart.parameter.TaskInstanceParameter` executes a task using the re
 
 
     class TaskB(gokart.TaskOnKart[str]):
-        require_task = gokart.TaskInstanceParameter()
+        require_task: gokart.TaskInstanceParameter = gokart.TaskInstanceParameter()
 
         def requires(self):
             return self.require_task
@@ -120,7 +120,7 @@ Example
             return cls(**json.loads(s))
 
     class DummyTask(gokart.TaskOnKart):
-        config: Config = gokart.SerializableParameter(object_type=Config)
+        config: gokart.SerializableParameter[Config] = gokart.SerializableParameter(object_type=Config)
 
         def run(self):
             # Save the `config` object as part of the task result.

docs/task_settings.rst

@@ -83,7 +83,7 @@ If you want to change the parameter of TaskA and rerun TaskB and TaskC, recommen
 .. code:: python
 
     class TaskA(gokart.TaskOnKart):
-        __version = luigi.IntParameter(default=1)
+        __version: luigi.IntParameter = luigi.IntParameter(default=1)
 
 If the hash value of TaskA will change, the dependent tasks (in this case, TaskB and TaskC) will rerun.
 

docs/tutorial.rst

@@ -147,7 +147,7 @@ Modify ``example/gokart_example/model/sample.py`` as follows:
 
     class StringToSplit(GokartTask):
         """Like the function to divide received data by spaces."""
-        task = gokart.TaskInstanceParameter()
+        task: gokart.TaskInstanceParameter = gokart.TaskInstanceParameter()
 
         def run(self):
             sample = self.load('task')
@@ -240,7 +240,7 @@ Add new parameter on dependent tasks like following:
 .. code:: python
 
     class Sample(GokartTask):
-        version = luigi.IntParameter(default=1)
+        version: luigi.IntParameter = luigi.IntParameter(default=1)
 
         def run(self):
             self.dump('sample output version {self.version}')

gokart/gcs_config.py

@@ -10,7 +10,7 @@
 
 
 class GCSConfig(luigi.Config):
-    gcs_credential_name: str = luigi.Parameter(default='GCS_CREDENTIAL', description='GCS credential environment variable.')
+    gcs_credential_name: luigi.StrParameter = luigi.StrParameter(default='GCS_CREDENTIAL', description='GCS credential environment variable.')
     _client = None
 
     def get_gcs_client(self) -> luigi.contrib.gcs.GCSClient:

gokart/info.py

@@ -45,8 +45,8 @@ def make_tree_info(
 
 
 class tree_info(TaskOnKart[Any]):
-    mode: str = luigi.Parameter(default='', description='This must be in ["simple", "all"].')
-    output_path: str = luigi.Parameter(default='tree.txt', description='Output file path.')
+    mode: luigi.StrParameter = luigi.StrParameter(default='', description='This must be in ["simple", "all"].')
+    output_path: luigi.StrParameter = luigi.StrParameter(default='tree.txt', description='Output file path.')
 
     def output(self):
         return self.make_target(self.output_path, use_unique_id=False)

gokart/parameter.py

@@ -3,27 +3,48 @@
 import bz2
 import datetime
 import json
+import sys
 from logging import getLogger
 from typing import Any, Generic, Protocol, TypeVar
+
+if sys.version_info >= (3, 11):
+    from typing import Unpack
+else:
+    from typing_extensions import Unpack
 from warnings import warn
 
 import luigi
 from luigi import task_register
 
+try:
+    from luigi.parameter import _no_value, _NoValueType, _ParameterKwargs
+except ImportError:
+    _no_value = None  # type: ignore[assignment]
+    _NoValueType = type(None)  # type: ignore[assignment,misc]
+    _ParameterKwargs = dict  # type: ignore[assignment,misc]
+
 import gokart
 
 logger = getLogger(__name__)
 
 
-class TaskInstanceParameter(luigi.Parameter):
-    def __init__(self, expected_type=None, *args, **kwargs):
+TASK_ON_KART_TYPE = TypeVar('TASK_ON_KART_TYPE', bound='gokart.TaskOnKart')  # type: ignore
+
+
+class TaskInstanceParameter(luigi.Parameter[TASK_ON_KART_TYPE], Generic[TASK_ON_KART_TYPE]):
+    def __init__(
+        self,
+        expected_type: type[TASK_ON_KART_TYPE] | None = None,
+        default: TASK_ON_KART_TYPE | _NoValueType = _no_value,
+        **kwargs: Unpack[_ParameterKwargs],
+    ):
         if expected_type is None:
             self.expected_type: type = gokart.TaskOnKart
         elif isinstance(expected_type, type):
             self.expected_type = expected_type
         else:
             raise TypeError(f'expected_type must be a type, not {type(expected_type)}')
-        super().__init__(*args, **kwargs)
+        super().__init__(default=default, **kwargs)
 
     @staticmethod
     def _recursive(param_dict):
@@ -64,15 +85,20 @@ def default(self, obj):
         return json.JSONEncoder.default(self, obj)
 
 
-class ListTaskInstanceParameter(luigi.Parameter):
-    def __init__(self, expected_elements_type=None, *args, **kwargs):
+class ListTaskInstanceParameter(luigi.Parameter[list[TASK_ON_KART_TYPE]], Generic[TASK_ON_KART_TYPE]):
+    def __init__(
+        self,
+        expected_elements_type: type[TASK_ON_KART_TYPE] | None = None,
+        default: list[TASK_ON_KART_TYPE] | _NoValueType = _no_value,
+        **kwargs: Unpack[_ParameterKwargs],
+    ):
         if expected_elements_type is None:
             self.expected_elements_type: type = gokart.TaskOnKart
         elif isinstance(expected_elements_type, type):
             self.expected_elements_type = expected_elements_type
         else:
             raise TypeError(f'expected_elements_type must be a type, not {type(expected_elements_type)}')
-        super().__init__(*args, **kwargs)
+        super().__init__(default=default, **kwargs)
 
     def parse(self, s):
         return [TaskInstanceParameter().parse(x) for x in list(json.loads(s))]
@@ -113,7 +139,7 @@ def gokart_deserialize(cls: type[T], s: str) -> T:
 S = TypeVar('S', bound=Serializable)
 
 
-class SerializableParameter(luigi.Parameter, Generic[S]):
+class SerializableParameter(luigi.Parameter[S], Generic[S]):
     def __init__(self, object_type: type[S], *args: Any, **kwargs: Any) -> None:
         self._object_type = object_type
         super().__init__(*args, **kwargs)
@@ -125,7 +151,7 @@ def serialize(self, x: S) -> str:
         return x.gokart_serialize()
 
 
-class ZonedDateSecondParameter(luigi.Parameter):
+class ZonedDateSecondParameter(luigi.Parameter[datetime.datetime]):
     """
     ZonedDateSecondParameter supports a datetime.datetime object with timezone information.