Merge pull request #150 from nzw0301/ruff

Use `__future__.annotations`
optuna · Aug 20, 2024 · 605390c · 605390c
2 parents f9a71d3 + af3e5fb
commit 605390c
Show file tree

Hide file tree

Showing 35 changed files with 220 additions and 193 deletions.
diff --git a/optuna_integration/_imports.py b/optuna_integration/_imports.py
@@ -1,10 +1,9 @@
+from __future__ import annotations
+
 import importlib
 import types
 from types import TracebackType
 from typing import Any
-from typing import Optional
-from typing import Tuple
-from typing import Type
 
 
 class _DeferredImportExceptionContextManager:
@@ -16,7 +15,7 @@ class _DeferredImportExceptionContextManager:
     """
 
     def __init__(self) -> None:
-        self._deferred: Optional[Tuple[Exception, str]] = None
+        self._deferred: tuple[Exception, str] | None = None
 
     def __enter__(self) -> "_DeferredImportExceptionContextManager":
         """Enter the context manager.
@@ -29,10 +28,10 @@ def __enter__(self) -> "_DeferredImportExceptionContextManager":
 
     def __exit__(
         self,
-        exc_type: Optional[Type[Exception]],
-        exc_value: Optional[Exception],
-        traceback: Optional[TracebackType],
-    ) -> Optional[bool]:
+        exc_type: type[Exception] | None,
+        exc_value: Exception | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
         """Exit the context manager.
 
         Args:

diff --git a/optuna_integration/_lightgbm_tuner/sklearn.py b/optuna_integration/_lightgbm_tuner/sklearn.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any
 import warnings
 

diff --git a/optuna_integration/allennlp/_dump_best_config.py b/optuna_integration/allennlp/_dump_best_config.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 
 import optuna

diff --git a/optuna_integration/allennlp/_pruner.py b/optuna_integration/allennlp/_pruner.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 import os
 from typing import Any
-from typing import Callable
 
 from optuna import load_study
 from optuna import pruners

diff --git a/optuna_integration/allennlp/_variables.py b/optuna_integration/allennlp/_variables.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import os
 from typing import Any

diff --git a/optuna_integration/botorch/botorch.py b/optuna_integration/botorch/botorch.py
@@ -1,11 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Callable
+from collections.abc import Sequence
 from typing import Any
-from typing import Callable
-from typing import Dict
-from typing import Optional
-from typing import Sequence
-from typing import Union
 import warnings
 
 import numpy
@@ -100,9 +97,9 @@ def _get_constraint_funcs(n_constraints: int) -> list[Callable[["torch.Tensor"],
 def logei_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Log Expected Improvement (LogEI).
 
@@ -210,9 +207,9 @@ def logei_candidates_func(
 def qei_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Quasi MC-based batch Expected Improvement (qEI).
 
@@ -315,9 +312,9 @@ def qei_candidates_func(
 def qnei_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Quasi MC-based batch Noisy Expected Improvement (qNEI).
 
@@ -382,9 +379,9 @@ def qnei_candidates_func(
 def qehvi_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Quasi MC-based batch Expected Hypervolume Improvement (qEHVI).
 
@@ -467,9 +464,9 @@ def qehvi_candidates_func(
 def ehvi_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Expected Hypervolume Improvement (EHVI).
 
@@ -532,9 +529,9 @@ def ehvi_candidates_func(
 def qnehvi_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Quasi MC-based batch Noisy Expected Hypervolume Improvement (qNEHVI).
 
@@ -616,9 +613,9 @@ def qnehvi_candidates_func(
 def qparego_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Quasi MC-based extended ParEGO (qParEGO) for constrained multi-objective optimization.
 
@@ -688,9 +685,9 @@ def qparego_candidates_func(
 def qkg_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Quasi MC-based batch Knowledge Gradient (qKG).
 
@@ -754,9 +751,9 @@ def qkg_candidates_func(
 def qhvkg_candidates_func(
     train_x: "torch.Tensor",
     train_obj: "torch.Tensor",
-    train_con: Optional["torch.Tensor"],
+    train_con: "torch.Tensor" | None,
     bounds: "torch.Tensor",
-    pending_x: Optional["torch.Tensor"],
+    pending_x: "torch.Tensor" | None,
 ) -> "torch.Tensor":
     """Quasi MC-based batch Hypervolume Knowledge Gradient (qHVKG).
 
@@ -842,9 +839,9 @@ def _get_default_candidates_func(
     [
         "torch.Tensor",
         "torch.Tensor",
-        Optional["torch.Tensor"],
+        "torch.Tensor" | None,
         "torch.Tensor",
-        Optional["torch.Tensor"],
+        "torch.Tensor" | None,
     ],
     "torch.Tensor",
 ]:
@@ -936,24 +933,25 @@ class BoTorchSampler(BaseSampler):
     def __init__(
         self,
         *,
-        candidates_func: Optional[
+        candidates_func: (
             Callable[
                 [
                     "torch.Tensor",
                     "torch.Tensor",
-                    Optional["torch.Tensor"],
+                    "torch.Tensor" | None,
                     "torch.Tensor",
-                    Optional["torch.Tensor"],
+                    "torch.Tensor" | None,
                 ],
                 "torch.Tensor",
             ]
-        ] = None,
-        constraints_func: Optional[Callable[[FrozenTrial], Sequence[float]]] = None,
+            | None
+        ) = None,
+        constraints_func: Callable[[FrozenTrial], Sequence[float]] | None = None,
         n_startup_trials: int = 10,
         consider_running_trials: bool = False,
-        independent_sampler: Optional[BaseSampler] = None,
-        seed: Optional[int] = None,
-        device: Optional["torch.device"] = None,
+        independent_sampler: BaseSampler | None = None,
+        seed: int | None = None,
+        device: "torch.device" | None = None,
     ):
         _imports.check()
 
@@ -964,23 +962,23 @@ def __init__(
         self._n_startup_trials = n_startup_trials
         self._seed = seed
 
-        self._study_id: Optional[int] = None
+        self._study_id: int | None = None
         self._search_space = IntersectionSearchSpace()
         self._device = device or torch.device("cpu")
 
     def infer_relative_search_space(
         self,
         study: Study,
         trial: FrozenTrial,
-    ) -> Dict[str, BaseDistribution]:
+    ) -> dict[str, BaseDistribution]:
         if self._study_id is None:
             self._study_id = study._study_id
         if self._study_id != study._study_id:
             # Note that the check below is meaningless when `InMemoryStorage` is used
             # because `InMemoryStorage.create_new_study` always returns the same study ID.
             raise RuntimeError("BoTorchSampler cannot handle multiple studies.")
 
-        search_space: Dict[str, BaseDistribution] = {}
+        search_space: dict[str, BaseDistribution] = {}
         for name, distribution in self._search_space.calculate(study).items():
             if distribution.single():
                 # built-in `candidates_func` cannot handle distributions that contain just a
@@ -995,8 +993,8 @@ def sample_relative(
         self,
         study: Study,
         trial: FrozenTrial,
-        search_space: Dict[str, BaseDistribution],
-    ) -> Dict[str, Any]:
+        search_space: dict[str, BaseDistribution],
+    ) -> dict[str, Any]:
         assert isinstance(search_space, dict)
 
         if len(search_space) == 0:
@@ -1015,12 +1013,12 @@ def sample_relative(
 
         trans = _SearchSpaceTransform(search_space)
         n_objectives = len(study.directions)
-        values: Union[numpy.ndarray, torch.Tensor] = numpy.empty(
+        values: numpy.ndarray | torch.Tensor = numpy.empty(
             (n_trials, n_objectives), dtype=numpy.float64
         )
-        params: Union[numpy.ndarray, torch.Tensor]
-        con: Optional[Union[numpy.ndarray, torch.Tensor]] = None
-        bounds: Union[numpy.ndarray, torch.Tensor] = trans.bounds
+        params: numpy.ndarray | torch.Tensor
+        con: numpy.ndarray | torch.Tensor | None = None
+        bounds: numpy.ndarray | torch.Tensor = trans.bounds
         params = numpy.empty((n_trials, trans.bounds.shape[0]), dtype=numpy.float64)
         for trial_idx, trial in enumerate(trials):
             if trial.state == TrialState.COMPLETE:
@@ -1151,7 +1149,7 @@ def after_trial(
         study: Study,
         trial: FrozenTrial,
         state: TrialState,
-        values: Optional[Sequence[float]],
+        values: Sequence[float] | None,
     ) -> None:
         if self._constraints_func is not None:
             _process_constraints_after_trial(self._constraints_func, study, trial, state)

diff --git a/optuna_integration/chainermn/chainermn.py b/optuna_integration/chainermn/chainermn.py
@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+from collections.abc import Callable
+from collections.abc import Sequence
 from datetime import datetime
 from typing import Any
-from typing import Callable
 from typing import overload
-from typing import Sequence
 import warnings
 
 from optuna import TrialPruned