Switch to TypeAliasType where possible

ethicml/data/util.py

@@ -1,7 +1,7 @@
 """Useful methods that are used in some of the data objects."""
 
 from collections.abc import Mapping, Sequence
-from itertools import groupby
+from itertools import chain, groupby
 from typing import NamedTuple, TypeAlias
 
 import pandas as pd
@@ -98,7 +98,7 @@ def flatten_dict(dictionary: Mapping[str, list[str]] | None) -> list[str]:
     """Flatten a dictionary of lists by joining all lists to one big list."""
     if dictionary is None:
         return []
-    return [x for inner_list in dictionary.values() for x in inner_list]
+    return list(chain.from_iterable(dictionary.values()))
 
 
 def reduce_feature_group(

ethicml/models/inprocess/in_subprocess.py

@@ -7,8 +7,8 @@
 import json
 from pathlib import Path
 from tempfile import TemporaryDirectory, gettempdir
-from typing import Any, Literal, TypeAlias, TypedDict, final
-from typing_extensions import Self
+from typing import Any, Literal, TypedDict, final
+from typing_extensions import Self, TypeAliasType
 import uuid
 
 from ethicml.models.algorithm_base import SubprocessAlgorithmMixin
@@ -47,7 +47,7 @@ class InAlgoPredArgs(TypedDict):
     model: str
 
 
-InAlgoArgs: TypeAlias = InAlgoFitArgs | InAlgoPredArgs | InAlgoRunArgs
+InAlgoArgs = TypeAliasType("InAlgoArgs", InAlgoFitArgs | InAlgoPredArgs | InAlgoRunArgs)
 
 
 @dataclass

ethicml/models/inprocess/shared.py

@@ -1,6 +1,6 @@
 """Methods that are shared among the inprocess algorithms."""
 
-from typing import Union
+from typing_extensions import TypeAliasType
 
 from sklearn.linear_model import LogisticRegression
 from sklearn.svm import SVC, LinearSVC
@@ -9,8 +9,7 @@
 
 __all__ = ["LinearModel", "settings_for_svm_lr"]
 
-
-LinearModel = Union[SVC, LinearSVC, LogisticRegression]  # noqa: UP007
+LinearModel = TypeAliasType("LinearModel", SVC | LinearSVC | LogisticRegression)
 
 
 def settings_for_svm_lr(

ethicml/models/preprocess/beutel.py

@@ -55,9 +55,9 @@ def _get_flags(self) -> BeutelArgs:
         # TODO: replace this with dataclasses.asdict()
         return {
             "fairness": self.fairness,
-            "enc_size": list(self.enc_size),
-            "adv_size": list(self.adv_size),
-            "pred_size": list(self.pred_size),
+            "enc_size": self.enc_size,
+            "adv_size": self.adv_size,
+            "pred_size": self.pred_size,
             "enc_activation": self.enc_activation,
             "adv_activation": self.adv_activation,
             "batch_size": self.batch_size,

ethicml/models/preprocess/pre_subprocess.py

@@ -6,8 +6,8 @@
 import json
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Any, Literal, TypeAlias, TypedDict, TypeVar, final
-from typing_extensions import Self
+from typing import Any, Literal, TypedDict, TypeVar, final
+from typing_extensions import Self, TypeAliasType
 
 from ethicml.models.algorithm_base import SubprocessAlgorithmMixin
 from ethicml.models.preprocess.pre_algorithm import PreAlgorithm
@@ -49,7 +49,7 @@ class PreAlgoTformArgs(TypedDict):
     model: str
 
 
-PreAlgoArgs: TypeAlias = PreAlgoFitArgs | PreAlgoTformArgs | PreAlgoRunArgs
+PreAlgoArgs = TypeAliasType("PreAlgoArgs", PreAlgoFitArgs | PreAlgoTformArgs | PreAlgoRunArgs)
 
 
 @dataclass

ethicml/plot/common.py

@@ -1,13 +1,14 @@
 """Common plotting functions / datastructures."""
 
-from typing import Any, Literal, NamedTuple, TypeAlias
+from typing import Any, Literal, NamedTuple
+from typing_extensions import TypeAliasType
 
 from matplotlib import legend
 from matplotlib.axes import Axes
 import pandas as pd
 
-LegendType: TypeAlias = Literal["inside", "outside"]
-PlotType: TypeAlias = Literal["box", "cross", "scatter", "line"]
+LegendType = TypeAliasType("LegendType", Literal["inside", "outside"])
+PlotType = TypeAliasType("PlotType", Literal["box", "cross", "scatter", "line"])
 
 
 class DataEntry(NamedTuple):

ethicml/run/parallelism.py

@@ -1,7 +1,8 @@
 """Collection of functions that enable parallelism."""
 
 from collections.abc import Sequence
-from typing import Protocol, TypeAlias, TypeVar, cast, overload
+from typing import Protocol, TypeVar, cast, overload
+from typing_extensions import TypeAliasType
 
 from joblib import Parallel, delayed
 import numpy as np
@@ -13,11 +14,11 @@
 
 __all__ = ["arrange_in_parallel", "run_in_parallel"]
 
-InSeq: TypeAlias = Sequence[InAlgorithm]
-PreSeq: TypeAlias = Sequence[PreAlgorithm]
-InResult: TypeAlias = list[list[Prediction]]
-PreResult: TypeAlias = list[list[tuple[DataTuple, DataTuple]]]
-DataSeq: TypeAlias = Sequence[TrainValPair]
+InSeq = TypeAliasType("InSeq", Sequence[InAlgorithm])
+PreSeq = TypeAliasType("PreSeq", Sequence[PreAlgorithm])
+InResult = TypeAliasType("InResult", list[list[Prediction]])
+PreResult = TypeAliasType("PreResult", list[list[tuple[DataTuple, DataTuple]]])
+DataSeq = TypeAliasType("DataSeq", Sequence[TrainValPair])
 
 
 @overload

ethicml/utility/data_structures.py

@@ -7,7 +7,7 @@
 import json
 from pathlib import Path
 from typing import Final, Literal, NamedTuple, NewType, TypeAlias, TypeVar, final
-from typing_extensions import Self
+from typing_extensions import Self, TypeAliasType
 
 import numpy as np
 from numpy import typing as npt
@@ -40,7 +40,7 @@
     "map_over_results_index",
 ]
 
-AxisType: TypeAlias = Literal["columns", "index"]
+AxisType = TypeAliasType("AxisType", Literal["columns", "index"])
 
 
 class PandasIndex(Enum):
@@ -371,8 +371,8 @@ def replace_data(self, data: pd.DataFrame, name: str | None = None) -> Self:
 EvalTuple: TypeAlias = LabelTuple | DataTuple
 """Union of :class:`LabelTuple` and :class:`DataTuple`."""
 
-HyperParamValue: TypeAlias = bool | int | float | str
-HyperParamType: TypeAlias = dict[str, HyperParamValue]
+HyperParamValue = TypeAliasType("HyperParamValue", bool | int | float | str)
+HyperParamType = TypeAliasType("HyperParamType", dict[str, HyperParamValue])
 
 T = TypeVar("T", SubgroupTuple, DataTuple)
 

typings/joblib/__init__.pyi

@@ -1,10 +1,10 @@
 from collections.abc import Callable, Iterable
 from pathlib import Path
 from types import TracebackType
-from typing import Any, Generic, Literal, TypeAlias, TypeVar, overload
-from typing_extensions import ParamSpec, Self
+from typing import Any, Generic, Literal, TypeVar, overload
+from typing_extensions import ParamSpec, Self, TypeAliasType
 
-_CompressLevel: TypeAlias = Literal[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+_CompressLevel = TypeAliasType("_CompressLevel", Literal[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
 
 def dump(
     value: Any,