Fix 101 mypy issues

Author: eddiebergman

autosklearn/automl.py

@@ -1,5 +1,4 @@
-# -*- encoding: utf-8 -*-
-from typing import List, Optional, Tuple, Union
+from typing import overload
 
 import logging
 
@@ -16,7 +15,7 @@
 
 def convert_if_sparse(
     y: SUPPORTED_TARGET_TYPES,
-) -> Union[np.ndarray, List, pd.DataFrame, pd.Series]:
+) -> np.ndarray | list | pd.DataFrame | pd.Series:
     """If the labels `y` are sparse, it will convert it to its dense representation
 
     Parameters
@@ -77,9 +76,9 @@ class InputValidator(BaseEstimator):
 
     def __init__(
         self,
-        feat_type: Optional[List[str]] = None,
+        feat_type: list[str] | None = None,
         is_classification: bool = False,
-        logger_port: Optional[int] = None,
+        logger_port: int | None = None,
         allow_string_features: bool = True,
     ) -> None:
         self.feat_type = feat_type
@@ -108,8 +107,8 @@ def fit(
         self,
         X_train: SUPPORTED_FEAT_TYPES,
         y_train: SUPPORTED_TARGET_TYPES,
-        X_test: Optional[SUPPORTED_FEAT_TYPES] = None,
-        y_test: Optional[SUPPORTED_TARGET_TYPES] = None,
+        X_test: SUPPORTED_FEAT_TYPES | None = None,
+        y_test: SUPPORTED_TARGET_TYPES | None = None,
     ) -> BaseEstimator:
         """
         Validates and fit a categorical encoder (if needed) to the features, and
@@ -172,11 +171,27 @@ def fit(
 
         return self
 
+    @overload
     def transform(
         self,
         X: SUPPORTED_FEAT_TYPES,
-        y: Optional[Union[List, pd.Series, pd.DataFrame, np.ndarray]] = None,
-    ) -> Tuple[Union[np.ndarray, pd.DataFrame, spmatrix], Optional[np.ndarray]]:
+        y: None = None,
+    ) -> tuple[np.ndarray | pd.DataFrame | spmatrix, None]:
+        ...
+
+    @overload
+    def transform(
+        self,
+        X: SUPPORTED_FEAT_TYPES,
+        y: list | pd.Series | pd.DataFrame | np.ndarray,
+    ) -> tuple[np.ndarray | pd.DataFrame | spmatrix, np.ndarray]:
+        ...
+
+    def transform(
+        self,
+        X: SUPPORTED_FEAT_TYPES,
+        y: list | pd.Series | pd.DataFrame | np.ndarray | None = None,
+    ) -> tuple[np.ndarray | pd.DataFrame | spmatrix, np.ndarray | None]:
         """
         Transform the given target or features to a numpy array
 

autosklearn/data/validation.py

@@ -64,7 +64,7 @@ def __init__(
         else:
             raise ValueError(
                 "Unsupported feat_type provided. We expect the user to "
-                "provide a Dict[str, str] mapping from column to categorical/ "
+                "provide a Dict[str | int, str] mapping from column to categorical/ "
                 "numerical."
             )
 

autosklearn/data/xy_data_manager.py

@@ -1,4 +1,3 @@
-# -*- encoding: utf-8 -*-
 from __future__ import annotations
 
 from typing import (
@@ -12,6 +11,7 @@
     Type,
     Union,
     cast,
+    overload,
 )
 
 import functools
@@ -98,14 +98,27 @@ def fit_predict_try_except_decorator(
         queue.close()
 
 
-def get_cost_of_crash(metrics: Sequence[Scorer]) -> List[float] | float:
+@overload
+def get_cost_of_crash(metrics: Sequence[Scorer]) -> List[float]:
+    ...
+
+
+@overload
+def get_cost_of_crash(metrics: Scorer) -> float:
+    ...
+
+
+def get_cost_of_crash(metrics: Sequence[Scorer] | Scorer) -> List[float] | float:
     """Return the cost of crash.
 
     Return value can be either a list (multi-objective optimization) or a
     raw float (single objective) because SMAC assumes different types in the
     two different cases.
     """
     costs = []
+    if isinstance(metrics, Scorer):
+        metrics = [metrics]
+
     for metric in metrics:
         if not isinstance(metric, Scorer):
             raise ValueError("The metric {metric} must be an instance of Scorer")

autosklearn/evaluation/__init__.py

@@ -29,6 +29,7 @@
     SUPPORTED_TARGET_TYPES,
     SUPPORTED_XDATA_TYPES,
 )
+from autosklearn.data.validation import SUPPORTED_FEAT_TYPES
 from autosklearn.metrics import Scorer, calculate_losses
 from autosklearn.pipeline.components.base import ThirdPartyComponents, _addons
 from autosklearn.pipeline.implementations.util import (
@@ -164,8 +165,8 @@ def get_additional_run_info(self) -> Optional[TYPE_ADDITIONAL_INFO]:
 def _fit_and_suppress_warnings(
     logger: Union[logging.Logger, PicklableClientLogger],
     model: BaseEstimator,
-    X: np.ndarray,
-    y: np.ndarray,
+    X: SUPPORTED_FEAT_TYPES,
+    y: SUPPORTED_TARGET_TYPES,
 ) -> BaseEstimator:
     def send_warnings_to_log(
         message: Union[Warning, str],

autosklearn/evaluation/abstract_evaluator.py

@@ -115,7 +115,7 @@ def _fit_with_budget(
     budget_type: Optional[str],
     logger: Union[logging.Logger, PicklableClientLogger],
     model: BaseEstimator,
-    train_indices: List[int],
+    train_indices: Sequence[int] | np.ndarray,
     task_type: int,
 ) -> None:
     if (

autosklearn/evaluation/train_evaluator.py

@@ -431,7 +431,7 @@ def make_scorer(
 
 def _validate_metrics(
     metrics: Sequence[Scorer],
-    scoring_functions: Optional[List[Scorer]] = None,
+    scoring_functions: Optional[Sequence[Scorer]] = None,
 ) -> None:
     """
     Validate metrics given to Auto-sklearn. Raises an Exception in case of a problem.

autosklearn/metrics/__init__.py

@@ -338,7 +338,7 @@ def get_estimator(self):
         return self.estimator
 
 
-class AutoSklearnChoice(object):
+class AutoSklearnChoice:
     def __init__(self, dataset_properties, random_state=None):
         """
         Parameters
@@ -370,6 +370,7 @@ def __init__(self, dataset_properties, random_state=None):
         # self.set_hyperparameters(self.configuration)
         self.choice = None
 
+    @classmethod
     def get_components(cls):
         raise NotImplementedError()
 

autosklearn/pipeline/components/base.py

@@ -30,7 +30,7 @@ def add_ohe(ohe: "OHEChoice") -> None:
 
 class OHEChoice(AutoSklearnChoice):
     @classmethod
-    def get_components(cls: BaseEstimator) -> Dict[str, BaseEstimator]:
+    def get_components(cls) -> Dict[str, BaseEstimator]:
         components: Dict[str, BaseEstimator] = OrderedDict()
         components.update(_ohes)
         components.update(additional_components.components)

autosklearn/pipeline/components/data_preprocessing/categorical_encoding/__init__.py

@@ -30,7 +30,7 @@ def add_mc(mc: BaseEstimator) -> None:
 
 class CoalescenseChoice(AutoSklearnChoice):
     @classmethod
-    def get_components(cls: BaseEstimator) -> Dict[str, BaseEstimator]:
+    def get_components(cls) -> Dict[str, BaseEstimator]:
         components: Dict[str, BaseEstimator] = OrderedDict()
         components.update(_mcs)
         components.update(additional_components.components)