Merge branch 'master' into introduce-polars

Author: hirosassa

docs/conf.py

@@ -1,5 +1,3 @@
-# -*- coding: utf-8 -*-
-
 # https://github.com/sphinx-doc/sphinx/issues/6211
 import luigi
 

gokart/build.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
+
 import enum
 import logging
 import sys
 from dataclasses import dataclass
 from functools import partial
 from logging import getLogger
-from typing import Literal, Optional, Protocol, TypeVar, cast, overload
+from typing import Literal, Protocol, TypeVar, cast, overload
 
 import backoff
 import luigi
@@ -62,7 +64,7 @@ def add(self, task: TaskOnKart) -> bool: ...
 
     def run(self) -> bool: ...
 
-    def __enter__(self) -> 'WorkerProtocol': ...
+    def __enter__(self) -> WorkerProtocol: ...
 
     def __exit__(self, type, value, traceback) -> Literal[False]: ...
 
@@ -162,7 +164,7 @@ def build(
     task_lock_exception_max_wait_seconds: int = 600,
     task_dump_config: TaskDumpConfig = TaskDumpConfig(),
     **env_params,
-) -> Optional[T]:
+) -> T | None:
     """
     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/build_process_task_info.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 
 import gokart

gokart/config_params.py

@@ -1,12 +1,12 @@
-from typing import Dict, Optional, Type
+from __future__ import annotations
 
 import luigi
 
 import gokart
 
 
 class inherits_config_params:
-    def __init__(self, config_class: Type[luigi.Config], parameter_alias: Optional[Dict[str, str]] = None):
+    def __init__(self, config_class: type[luigi.Config], parameter_alias: dict[str, str] | None = None):
         """
         Decorates task to inherit parameter value of `config_class`.
 
@@ -15,10 +15,10 @@ def __init__(self, config_class: Type[luigi.Config], parameter_alias: Optional[D
                            key: config_class's parameter name. value: decorated task's parameter name.
         """
 
-        self._config_class: Type[luigi.Config] = config_class
-        self._parameter_alias: Dict[str, str] = parameter_alias if parameter_alias is not None else {}
+        self._config_class: type[luigi.Config] = config_class
+        self._parameter_alias: dict[str, str] = parameter_alias if parameter_alias is not None else {}
 
-    def __call__(self, task_class: Type[gokart.TaskOnKart]):
+    def __call__(self, task_class: type[gokart.TaskOnKart]):
         # wrap task to prevent task name from being changed
         @luigi.task._task_wraps(task_class)
         class Wrapped(task_class):  # type: ignore
@@ -29,6 +29,6 @@ def get_param_values(cls, params, args, kwargs):
 
                     if hasattr(cls, task_param_key) and task_param_key not in kwargs:
                         kwargs[task_param_key] = param_value
-                return super(Wrapped, cls).get_param_values(params, args, kwargs)
+                return super().get_param_values(params, args, kwargs)
 
         return Wrapped

gokart/conflict_prevention_lock/task_lock.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import functools
 import os
 from logging import getLogger
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
 import redis
 from apscheduler.schedulers.background import BackgroundScheduler
@@ -10,9 +12,9 @@
 
 
 class TaskLockParams(NamedTuple):
-    redis_host: Optional[str]
-    redis_port: Optional[int]
-    redis_timeout: Optional[int]
+    redis_host: str | None
+    redis_port: int | None
+    redis_timeout: int | None
     redis_key: str
     should_task_lock: bool
     raise_task_lock_exception_on_collision: bool
@@ -31,10 +33,10 @@ def __new__(cls, *args, **kwargs):
         if cls not in cls._instances:
             cls._instances[cls] = {}
         if key not in cls._instances[cls]:
-            cls._instances[cls][key] = super(RedisClient, cls).__new__(cls)
+            cls._instances[cls][key] = super().__new__(cls)
         return cls._instances[cls][key]
 
-    def __init__(self, host: Optional[str], port: Optional[int]) -> None:
+    def __init__(self, host: str | None, port: int | None) -> None:
         if not hasattr(self, '_redis_client'):
             host = host or 'localhost'
             port = port or 6379
@@ -72,17 +74,17 @@ def set_lock_scheduler(task_lock: redis.lock.Lock, task_lock_params: TaskLockPar
     return scheduler
 
 
-def make_task_lock_key(file_path: str, unique_id: Optional[str]):
+def make_task_lock_key(file_path: str, unique_id: str | None):
     basename_without_ext = os.path.splitext(os.path.basename(file_path))[0]
     return f'{basename_without_ext}_{unique_id}'
 
 
 def make_task_lock_params(
     file_path: str,
-    unique_id: Optional[str],
-    redis_host: Optional[str] = None,
-    redis_port: Optional[int] = None,
-    redis_timeout: Optional[int] = None,
+    unique_id: str | None,
+    redis_host: str | None = None,
+    redis_port: int | None = None,
+    redis_timeout: int | None = None,
     raise_task_lock_exception_on_collision: bool = False,
     lock_extend_seconds: int = 10,
 ) -> TaskLockParams:

gokart/conflict_prevention_lock/task_lock_wrappers.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import functools
 from logging import getLogger
 from typing import Any, Callable

gokart/file_processor.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import xml.etree.ElementTree as ET
 from abc import abstractmethod
@@ -63,7 +65,7 @@ def dump(self, obj, file):
         file.write(obj)
 
 
-class _ChunkedLargeFileReader(object):
+class _ChunkedLargeFileReader:
     def __init__(self, file) -> None:
         self._file = file
 
@@ -183,8 +185,11 @@ def dump(self, obj, file):
 
 
 class JsonFileProcessor(FileProcessor):
+    def __init__(self, orient: str | None = None):
+        self._orient = orient
+
     def format(self):
-        return None
+        return luigi.format.Nop
 
     def load(self, file):
         ...
@@ -196,6 +201,8 @@ def dump(self, obj, file):
 class PolarsJsonFileProcessor(JsonFileProcessor):
     def load(self, file):
         try:
+            if self._orient == 'records':
+                return self.read_ndjson(file)
             return self.read_json(file)
         except pl.exceptions.NoDataError:
             return pl.DataFrame()
@@ -206,13 +213,17 @@ def dump(self, obj, file):
         )
         if isinstance(obj, dict):
             obj = pl.from_dict(obj)
-        obj.write_json(file)
+
+        if self._orient == 'records':
+            obj_write_ndjson(file)
+        else:
+            obj.write_json(file)
 
 
 class PandasJsonFileProcessor(JsonFileProcessor):
     def load(self, file):
         try:
-            return self.read_json(file)
+            return pd.read_json(file, orient=self._orient, lines=True if self._orient == 'records' else False)
         except pd.errors.EmptyDataError:
             return pd.DataFrame()
 
@@ -222,7 +233,7 @@ def dump(self, obj, file):
         )
         if isinstance(obj, dict):
             obj = pd.DataFrame.from_dict(obj)
-        obj.to_json(file)
+        obj.to_json(file, orient=self._orient, lines=True if self._orient == 'records' else False)
 
 
 class XmlFileProcessor(FileProcessor):
@@ -386,6 +397,7 @@ def make_file_processor(file_path: str, store_index_in_feather: bool) -> FilePro
         '.pkl': PickleFileProcessor(),
         '.gz': GzipFileProcessor(),
         '.json': JsonFileProcessor(),
+        '.ndjson': JsonFileProcessor(orient='records'),
         '.xml': XmlFileProcessor(),
         '.npz': NpzFileProcessor(),
         '.parquet': ParquetFileProcessor(compression='gzip'),

gokart/gcs_config.py

@@ -1,6 +1,7 @@
+from __future__ import annotations
+
 import json
 import os
-from typing import Optional
 
 import luigi
 import luigi.contrib.gcs
@@ -19,7 +20,7 @@ def get_gcs_client(self) -> luigi.contrib.gcs.GCSClient:
     def _get_gcs_client(self) -> luigi.contrib.gcs.GCSClient:
         return luigi.contrib.gcs.GCSClient(oauth_credentials=self._load_oauth_credentials())
 
-    def _load_oauth_credentials(self) -> Optional[Credentials]:
+    def _load_oauth_credentials(self) -> Credentials | None:
         json_str = os.environ.get(self.gcs_credential_name)
         if not json_str:
             return None

gokart/gcs_obj_metadata_client.py

@@ -3,7 +3,7 @@
 import copy
 import re
 from logging import getLogger
-from typing import Any, Union
+from typing import Any
 from urllib.parse import urlsplit
 
 from googleapiclient.model import makepatch
@@ -84,7 +84,7 @@ def _get_patched_obj_metadata(
         metadata: Any,
         task_params: dict[str, str] | None = None,
         custom_labels: dict[str, Any] | None = None,
-    ) -> Union[dict, Any]:
+    ) -> dict | Any:
         # If metadata from response when getting bucket and object information is not dictionary,
         # something wrong might be happened, so return original metadata, no patched.
         if not isinstance(metadata, dict):

gokart/gcs_zip_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import shutil