lint

Author: OriolAbril

src/arviz_base/base.py

@@ -3,9 +3,9 @@
 import importlib
 import re
 import warnings
-from collections.abc import Hashable, Iterable, Mapping
+from collections.abc import Callable, Hashable, Iterable, Mapping
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import numpy as np
 import xarray as xr
@@ -24,9 +24,9 @@
 def generate_dims_coords(
     shape: Iterable[int],
     var_name: Hashable,
-    dims: Optional[Iterable[Hashable]] = None,
-    coords: Optional[CoordSpec] = None,
-    index_origin: Optional[int] = None,
+    dims: Iterable[Hashable] | None = None,
+    coords: CoordSpec | None = None,
+    index_origin: int | None = None,
     skip_event_dims: bool = False,
     check_conventions: bool = True,
 ) -> tuple[list[Hashable], CoordSpec]:
@@ -216,12 +216,12 @@ def ndarray_to_dataarray(
 def dict_to_dataset(
     data: DictData,
     *,
-    attrs: Optional[Mapping[Any, Any]] = None,
-    inference_library: Optional[str] = None,
-    coords: Optional[CoordSpec] = None,
-    dims: Optional[DimSpec] = None,
-    sample_dims: Optional[Iterable[Hashable]] = None,
-    index_origin: Optional[int] = None,
+    attrs: Mapping[Any, Any] | None = None,
+    inference_library: str | None = None,
+    coords: CoordSpec | None = None,
+    dims: DimSpec | None = None,
+    sample_dims: Iterable[Hashable] | None = None,
+    index_origin: int | None = None,
     skip_event_dims: bool = False,
     check_conventions: bool = True,
 ):
@@ -377,15 +377,15 @@ class requires:  # pylint: disable=invalid-name
     See https://github.com/arviz-devs/arviz/pull/1504 for more discussion.
     """
 
-    def __init__(self, *props: Union[str, list[str]]) -> None:
-        self.props: tuple[Union[str, list[str]], ...] = props
+    def __init__(self, *props: str | list[str]) -> None:
+        self.props: tuple[str | list[str], ...] = props
 
     def __call__(
         self, func: Callable[[RequiresArgTypeT], RequiresReturnTypeT]
-    ) -> Callable[[RequiresArgTypeT], Optional[RequiresReturnTypeT]]:  # noqa: D202
+    ) -> Callable[[RequiresArgTypeT], RequiresReturnTypeT | None]:  # noqa: D202
         """Wrap the decorated function."""
 
-        def wrapped(cls: RequiresArgTypeT) -> Optional[RequiresReturnTypeT]:
+        def wrapped(cls: RequiresArgTypeT) -> RequiresReturnTypeT | None:
             """Return None if not all props are available."""
             for prop in self.props:
                 prop_list = [prop] if isinstance(prop, str) else prop

src/arviz_base/io_dict.py

@@ -1,6 +1,5 @@
 """Dictionary specific conversion code."""
 import warnings
-from typing import Optional
 
 from datatree import DataTree
 
@@ -14,7 +13,7 @@ def from_dict(
     name=None,
     sample_dims=None,
     save_warmup=None,
-    index_origin: Optional[int] = None,
+    index_origin: int | None = None,
     coords=None,
     dims=None,
     pred_dims=None,

src/arviz_base/labels.py

@@ -1,6 +1,5 @@
 # pylint: disable=unused-argument
 """Utilities to generate labels from xarray objects."""
-from typing import Union
 
 __all__ = [
     "mix_labellers",
@@ -92,7 +91,7 @@ def sel_to_str(self, sel: dict, isel: dict):
             )
         return ""
 
-    def var_name_to_str(self, var_name: Union[str, None]):  # pylint: disable=no-self-use
+    def var_name_to_str(self, var_name: str | None):  # pylint: disable=no-self-use
         """WIP."""
         return var_name
 
@@ -106,7 +105,7 @@ def model_name_to_str(self, model_name):  # pylint: disable=no-self-use
         """WIP."""
         return model_name
 
-    def make_label_vert(self, var_name: Union[str, None], sel: dict, isel: dict):
+    def make_label_vert(self, var_name: str | None, sel: dict, isel: dict):
         """WIP."""
         var_name_str = self.var_name_to_str(var_name)
         sel_str = self.sel_to_str(sel, isel)

src/arviz_base/rcparams.py

@@ -7,10 +7,9 @@
 import sys
 from collections.abc import MutableMapping
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, get_args
 
 import numpy as np
-from typing_extensions import get_args
 
 _log = logging.getLogger("arviz")
 
@@ -203,7 +202,7 @@ def validate_iterable(value):
             if allow_auto and value.lower() == "auto":
                 return "auto"
             value = tuple(v.strip("([ ])") for v in value.split(",") if v.strip())
-        if np.iterable(value) and not isinstance(value, (set, frozenset)):
+        if np.iterable(value) and not isinstance(value, set | frozenset):
             val = tuple(scalar_validator(v) for v in value)
             if length is not None and len(val) != length:
                 raise ValueError(f"Iterable must be of length: {length}")

src/arviz_base/testing.py

@@ -1,13 +1,12 @@
 """ArviZ testing utilities."""
 
-from typing import Union
 
 from datatree import DataTree
 
 
 def check_multiple_attrs(
     test_dict: dict[str, list[str]], parent: DataTree
-) -> list[Union[str, tuple[str, str]]]:
+) -> list[str | tuple[str, str]]:
     """Perform multiple hasattr checks on InferenceData objects.
 
     It is thought to first check if the parent object contains a given dataset,
@@ -50,7 +49,7 @@ def check_multiple_attrs(
     in ``sample_stats``, also against what was expected.
 
     """
-    failed_attrs: list[Union[str, tuple[str, str]]] = []
+    failed_attrs: list[str | tuple[str, str]] = []
     for dataset_name, attributes in test_dict.items():
         if dataset_name.startswith("~"):
             if hasattr(parent, dataset_name[1:]):

src/arviz_base/utils.py

@@ -32,7 +32,7 @@ def _var_names(var_names, data, filter_vars=None):
         )
 
     if var_names is not None:
-        if isinstance(data, (list, tuple)):
+        if isinstance(data, list | tuple):
             all_vars = []
             for dataset in data:
                 dataset_vars = list(dataset.data_vars)
@@ -156,7 +156,7 @@ def _get_coords(data, coords):
     data : Dataset or DataArray
         Return type is of the same type as the input
     """
-    if not isinstance(data, (list, tuple)):
+    if not isinstance(data, list | tuple):
         try:
             return data.sel(**coords)
 
@@ -170,7 +170,7 @@ def _get_coords(data, coords):
                 "Check that coords structure is correct and"
                 f" dimensions are valid. {err}"
             ) from err
-    if not isinstance(coords, (list, tuple)):
+    if not isinstance(coords, list | tuple):
         coords = [coords] * len(data)
     data_subset = []
     for idx, (datum, coords_dict) in enumerate(zip(data, coords)):