feat: implement ZonedDateSecondParameter (#443)

* implement timestamp parameter minimally

* apply formatter

* change normalize only to remove microsecond

* add tests

* implement with zoned date second

* remove column separator case

---------

Co-authored-by: Shuntaro Takahashi <shuntaro-takahashi@m3.com>

Author: stakahashy

gokart/__init__.py

@@ -1,7 +1,13 @@
 from gokart.build import WorkerSchedulerFactory, build  # noqa:F401
 from gokart.info import make_tree_info, tree_info  # noqa:F401
 from gokart.pandas_type_config import PandasTypeConfig  # noqa:F401
-from gokart.parameter import ExplicitBoolParameter, ListTaskInstanceParameter, SerializableParameter, TaskInstanceParameter  # noqa:F401
+from gokart.parameter import (  # noqa:F401
+    ExplicitBoolParameter,
+    ListTaskInstanceParameter,
+    SerializableParameter,
+    TaskInstanceParameter,
+    ZonedDateSecondParameter,
+)
 from gokart.run import run  # noqa:F401
 from gokart.task import TaskOnKart  # noqa:F401
 from gokart.testing import test_run  # noqa:F401

gokart/parameter.py

@@ -1,7 +1,9 @@
 import bz2
+import datetime
 import json
 from logging import getLogger
 from typing import Generic, Protocol, TypeVar
+from warnings import warn
 
 import luigi
 from luigi import task_register
@@ -119,3 +121,35 @@ def parse(self, s: str) -> S:
 
     def serialize(self, x: S) -> str:
         return x.gokart_serialize()
+
+
+class ZonedDateSecondParameter(luigi.Parameter):
+    """
+    ZonedDateSecondParameter supports a datetime.datetime object with timezone information.
+
+    A ZonedDateSecondParameter is a `ISO 8601 <http://en.wikipedia.org/wiki/ISO_8601>`_ formatted
+    date, time specified to the second and timezone. For example, ``2013-07-10T19:07:38+09:00`` specifies July 10, 2013 at
+    19:07:38 +09:00. The separator `:` can be omitted for Python3.11 and later.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def parse(self, s):
+        # special character 'Z' is replaced with '+00:00'
+        # because Python 3.11 and later support fromisoformat with Z at the end of the string.
+        if s.endswith('Z'):
+            s = s[:-1] + '+00:00'
+        dt = datetime.datetime.fromisoformat(s)
+        if dt.tzinfo is None:
+            warn('The input does not have timezone information. Please consider using luigi.DateSecondParameter instead.', stacklevel=1)
+        return dt
+
+    def serialize(self, dt):
+        return dt.isoformat()
+
+    def normalize(self, dt):
+        # override _DatetimeParameterBase.normalize to avoid do nothing to normalize except removing microsecond.
+        # microsecond is removed because the number of digits of microsecond is not fixed.
+        # See also luigi's implementation  https://github.com/spotify/luigi/blob/v3.6.0/luigi/parameter.py#L612
+        return dt.replace(microsecond=0)

test/test_zoned_date_second_parameter.py

@@ -0,0 +1,59 @@
+import datetime
+import unittest
+
+from luigi.cmdline_parser import CmdlineParser
+
+from gokart import TaskOnKart, ZonedDateSecondParameter
+
+
+class ZonedDateSecondParameterTaskWithoutDefault(TaskOnKart):
+    task_namespace = __name__
+    dt: datetime.datetime = ZonedDateSecondParameter()
+
+    def run(self):
+        self.dump(self.dt)
+
+
+class ZonedDateSecondParameterTaskWithDefault(TaskOnKart):
+    task_namespace = __name__
+    dt: datetime.datetime = ZonedDateSecondParameter(default=datetime.datetime(2025, 2, 21, 12, 0, 0, tzinfo=datetime.timezone(datetime.timedelta(hours=9))))
+
+    def run(self):
+        self.dump(self.dt)
+
+
+class ZonedDateSecondParameterTest(unittest.TestCase):
+    def setUp(self):
+        self.default_datetime = datetime.datetime(2025, 2, 21, 12, 0, 0, tzinfo=datetime.timezone(datetime.timedelta(hours=9)))
+        self.default_datetime_str = '2025-02-21T12:00:00+09:00'
+
+    def test_default(self):
+        with CmdlineParser.global_instance([f'{__name__}.ZonedDateSecondParameterTaskWithDefault']) as cp:
+            assert cp.get_task_obj().dt == self.default_datetime
+
+    def test_parse_param_with_tz_suffix(self):
+        with CmdlineParser.global_instance([f'{__name__}.ZonedDateSecondParameterTaskWithDefault', '--dt', '2024-01-20T11:00:00+09:00']) as cp:
+            assert cp.get_task_obj().dt == datetime.datetime(2024, 1, 20, 11, 0, 0, tzinfo=datetime.timezone(datetime.timedelta(hours=9)))
+
+    def test_parse_param_with_Z_suffix(self):
+        with CmdlineParser.global_instance([f'{__name__}.ZonedDateSecondParameterTaskWithDefault', '--dt', '2024-01-20T11:00:00Z']) as cp:
+            assert cp.get_task_obj().dt == datetime.datetime(2024, 1, 20, 11, 0, 0, tzinfo=datetime.timezone(datetime.timedelta(hours=0)))
+
+    def test_parse_param_without_timezone_input(self):
+        with CmdlineParser.global_instance([f'{__name__}.ZonedDateSecondParameterTaskWithoutDefault', '--dt', '2025-02-21T12:00:00']) as cp:
+            assert cp.get_task_obj().dt == datetime.datetime(2025, 2, 21, 12, 0, 0, tzinfo=None)
+
+    def test_parse_method(self):
+        actual = ZonedDateSecondParameter().parse(self.default_datetime_str)
+        expected = self.default_datetime
+        self.assertEqual(actual, expected)
+
+    def test_serialize_task(self):
+        task = ZonedDateSecondParameterTaskWithoutDefault(dt=self.default_datetime)
+        actual = str(task)
+        expected = f'(dt={self.default_datetime_str})'
+        self.assertTrue(actual.endswith(expected))
+
+
+if __name__ == '__main__':
+    unittest.main()