Skip to content

Enforce ruff/flake8-pyi rules (PYI) #10359

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
May 29, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,7 @@ extend-select = [
"ISC", # flake8-implicit-str-concat
"PIE", # flake8-pie
"TID", # flake8-tidy-imports (absolute imports)
"PYI", # flake8-pyi
"I", # isort
"PERF", # Perflint
"W", # pycodestyle warnings
Expand All @@ -268,6 +269,8 @@ extend-safe-fixes = [
ignore = [
"C40", # unnecessary generator, comprehension, or literal
"PIE790", # unnecessary pass statement
"PYI019", # use `Self` instead of custom TypeVar
"PYI041", # use `float` instead of `int | float`
"PERF203", # try-except within a loop incurs performance overhead
"E402", # module level import not at top of file
"E731", # do not assign a lambda expression, use a def
Expand All @@ -283,6 +286,9 @@ ignore = [
[tool.ruff.lint.per-file-ignores]
# don't enforce absolute imports
"asv_bench/**" = ["TID252"]
# looks like ruff bugs
"xarray/core/_typed_ops.py" = ["PYI034"]
"xarray/namedarray/_typing.py" = ["PYI018", "PYI046"]

[tool.ruff.lint.isort]
known-first-party = ["xarray"]
Expand Down
10 changes: 5 additions & 5 deletions xarray/backends/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@
T_NetcdfEngine = Literal["netcdf4", "scipy", "h5netcdf"]
T_Engine = Union[
T_NetcdfEngine,
Literal["pydap", "zarr"],
Literal["pydap", "zarr"], # noqa: PYI051
type[BackendEntrypoint],
str, # no nice typing support for custom backends
None,
Expand Down Expand Up @@ -710,8 +710,8 @@ def open_dataset(
def open_dataarray(
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
*,
engine: T_Engine | None = None,
chunks: T_Chunks | None = None,
engine: T_Engine = None,
chunks: T_Chunks = None,
cache: bool | None = None,
decode_cf: bool | None = None,
mask_and_scale: bool | None = None,
Expand Down Expand Up @@ -1394,7 +1394,7 @@ def open_mfdataset(
| os.PathLike
| ReadBuffer
| NestedSequence[str | os.PathLike | ReadBuffer],
chunks: T_Chunks | None = None,
chunks: T_Chunks = None,
concat_dim: (
str
| DataArray
Expand All @@ -1406,7 +1406,7 @@ def open_mfdataset(
) = None,
compat: CompatOptions = "no_conflicts",
preprocess: Callable[[Dataset], Dataset] | None = None,
engine: T_Engine | None = None,
engine: T_Engine = None,
data_vars: Literal["all", "minimal", "different"] | list[str] = "all",
coords="different",
combine: Literal["by_coords", "nested"] = "by_coords",
Expand Down
16 changes: 7 additions & 9 deletions xarray/computation/apply_ufunc.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,16 +16,14 @@
Iterator,
Mapping,
Sequence,
Set,
)
from typing import TYPE_CHECKING, Any, Literal, TypeVar, Union
from collections.abc import (
Set as AbstractSet,
)
from typing import TYPE_CHECKING, Any, Literal

import numpy as np

_T = TypeVar("_T", bound=Union["Dataset", "DataArray"])
_U = TypeVar("_U", bound=Union["Dataset", "DataArray"])
_V = TypeVar("_V", bound=Union["Dataset", "DataArray"])

from xarray.core import duck_array_ops, utils
from xarray.core.formatting import limit_lines
from xarray.core.indexes import Index, filter_indexes_from_coords
Expand Down Expand Up @@ -200,7 +198,7 @@ def _get_coords_list(args: Iterable[Any]) -> list[Coordinates]:
def build_output_coords_and_indexes(
args: Iterable[Any],
signature: _UFuncSignature,
exclude_dims: Set = frozenset(),
exclude_dims: AbstractSet = frozenset(),
combine_attrs: CombineAttrsOptions = "override",
) -> tuple[list[dict[Any, Variable]], list[dict[Any, Index]]]:
"""Build output coordinates and indexes for an operation.
Expand Down Expand Up @@ -616,7 +614,7 @@ def apply_groupby_func(func, *args):


def unified_dim_sizes(
variables: Iterable[Variable], exclude_dims: Set = frozenset()
variables: Iterable[Variable], exclude_dims: AbstractSet = frozenset()
) -> dict[Hashable, int]:
dim_sizes: dict[Hashable, int] = {}

Expand Down Expand Up @@ -896,7 +894,7 @@ def apply_ufunc(
*args: Any,
input_core_dims: Sequence[Sequence] | None = None,
output_core_dims: Sequence[Sequence] | None = ((),),
exclude_dims: Set = frozenset(),
exclude_dims: AbstractSet = frozenset(),
vectorize: bool = False,
join: JoinOptions = "exact",
dataset_join: str = "exact",
Expand Down
3 changes: 2 additions & 1 deletion xarray/core/accessor_dt.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import numpy as np
import pandas as pd
from typing_extensions import Self

from xarray.coding.calendar_ops import _decimal_year
from xarray.coding.times import infer_calendar_name
Expand Down Expand Up @@ -650,7 +651,7 @@ def total_seconds(self) -> T_DataArray:
class CombinedDatetimelikeAccessor(
DatetimeAccessor[T_DataArray], TimedeltaAccessor[T_DataArray]
):
def __new__(cls, obj: T_DataArray) -> CombinedDatetimelikeAccessor:
def __new__(cls, obj: T_DataArray) -> Self:
# CombinedDatetimelikeAccessor isn't really instantiated. Instead
# we need to choose which parent (datetime or timedelta) is
# appropriate. Since we're checking the dtypes anyway, we'll just
Expand Down
9 changes: 6 additions & 3 deletions xarray/core/datatree_render.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,18 @@

from __future__ import annotations

from collections import namedtuple
from collections.abc import Iterable, Iterator
from math import ceil
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, NamedTuple

if TYPE_CHECKING:
from xarray.core.datatree import DataTree

Row = namedtuple("Row", ("pre", "fill", "node"))

class Row(NamedTuple):
pre: str
fill: str
node: DataTree | str


class AbstractStyle:
Expand Down
4 changes: 2 additions & 2 deletions xarray/core/indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1657,7 +1657,7 @@ class Indexes(collections.abc.Mapping, Generic[T_PandasOrXarrayIndex]):

"""

_index_type: type[Index] | type[pd.Index]
_index_type: type[Index | pd.Index]
_indexes: dict[Any, T_PandasOrXarrayIndex]
_variables: dict[Any, Variable]

Expand All @@ -1675,7 +1675,7 @@ def __init__(
self,
indexes: Mapping[Any, T_PandasOrXarrayIndex] | None = None,
variables: Mapping[Any, Variable] | None = None,
index_type: type[Index] | type[pd.Index] = Index,
index_type: type[Index | pd.Index] = Index,
):
"""Constructor not for public consumption.

Expand Down
6 changes: 3 additions & 3 deletions xarray/core/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ def copy(

# FYI in some cases we don't allow `None`, which this doesn't take account of.
# FYI the `str` is for a size string, e.g. "16MB", supported by dask.
T_ChunkDim: TypeAlias = str | int | Literal["auto"] | None | tuple[int, ...]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the intent was to have int, tuple of int or "auto" — not any string. I hope I'm not wrong.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

check comment above :)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was so busy scavenging the reference documentation for a hint that I didn't see the comment 😕

Reverted and added # noqa: PYI051 instead. I could remove Literal["auto"] since it is covered by str, but I'd rather keep it for code documentation purposes.

T_ChunkDim: TypeAlias = str | int | Literal["auto"] | None | tuple[int, ...] # noqa: PYI051
T_ChunkDimFreq: TypeAlias = Union["TimeResampler", T_ChunkDim]
T_ChunksFreq: TypeAlias = T_ChunkDim | Mapping[Any, T_ChunkDimFreq]
# We allow the tuple form of this (though arguably we could transition to named dims only)
Expand Down Expand Up @@ -329,7 +329,7 @@ def mode(self) -> str:
# for _get_filepath_or_buffer
...

def seek(self, __offset: int, __whence: int = ...) -> int:
def seek(self, offset: int, whence: int = ..., /) -> int:
# with one argument: gzip.GzipFile, bz2.BZ2File
# with two arguments: zip.ZipFile, read_sas
...
Expand All @@ -345,7 +345,7 @@ def tell(self) -> int:

@runtime_checkable
class ReadBuffer(BaseBuffer, Protocol[AnyStr_co]):
def read(self, __n: int = ...) -> AnyStr_co:
def read(self, n: int = ..., /) -> AnyStr_co:
# for BytesIOWrapper, gzip.GzipFile, bz2.BZ2File
...

Expand Down
6 changes: 4 additions & 2 deletions xarray/core/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,11 @@
MutableMapping,
MutableSet,
Sequence,
Set,
ValuesView,
)
from collections.abc import (
Set as AbstractSet,
)
from enum import Enum
from pathlib import Path
from types import EllipsisType, ModuleType
Expand Down Expand Up @@ -1055,7 +1057,7 @@ def parse_ordered_dims(
)


def _check_dims(dim: Set[Hashable], all_dims: Set[Hashable]) -> None:
def _check_dims(dim: AbstractSet[Hashable], all_dims: AbstractSet[Hashable]) -> None:
wrong_dims = (dim - all_dims) - {...}
if wrong_dims:
wrong_dims_str = ", ".join(f"'{d}'" for d in wrong_dims)
Expand Down
2 changes: 1 addition & 1 deletion xarray/groupers.py
Original file line number Diff line number Diff line change
Expand Up @@ -902,7 +902,7 @@ def factorize(self, group: T_Group) -> EncodedGroups:
first_items = agged["first"]
counts = agged["count"]

index_class: type[CFTimeIndex] | type[pd.DatetimeIndex]
index_class: type[CFTimeIndex | pd.DatetimeIndex]
if _contains_cftime_datetimes(group.data):
index_class = CFTimeIndex
datetime_class = type(first_n_items(group.data, 1).item())
Expand Down
3 changes: 1 addition & 2 deletions xarray/namedarray/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@ class Default(Enum):
_default = Default.token

# https://stackoverflow.com/questions/74633074/how-to-type-hint-a-generic-numpy-array
_T = TypeVar("_T")
Copy link
Contributor Author

@DimitriPapadopoulos DimitriPapadopoulos May 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we keep this (private) type even though it is not used?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we stopped using it, unless the link in the line above is relevant?

Copy link
Contributor Author

@DimitriPapadopoulos DimitriPapadopoulos May 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I may be mistaken, but https://stackoverflow.com/questions/74633074/how-to-type-hint-a-generic-numpy-array does not seem relevant to _T vs. _T_co. Since both _typing and _T are private, I guess it's safe to remove _T.

Copy link
Contributor Author

@DimitriPapadopoulos DimitriPapadopoulos May 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's defined locally elsewhere:

_T = TypeVar("_T")

It looks like the more global definition in xarray/namedarray/_typing.py is not used any more.

_T_co = TypeVar("_T_co", covariant=True)

_dtype = np.dtype
Expand Down Expand Up @@ -79,7 +78,7 @@ def dtype(self) -> _DType_co: ...
_NormalizedChunks = tuple[tuple[int, ...], ...]
# FYI in some cases we don't allow `None`, which this doesn't take account of.
# # FYI the `str` is for a size string, e.g. "16MB", supported by dask.
T_ChunkDim: TypeAlias = str | int | Literal["auto"] | None | tuple[int, ...]
T_ChunkDim: TypeAlias = str | int | Literal["auto"] | None | tuple[int, ...] # noqa: PYI051
# We allow the tuple form of this (though arguably we could transition to named dims only)
T_Chunks: TypeAlias = T_ChunkDim | Mapping[Any, T_ChunkDim]

Expand Down
8 changes: 4 additions & 4 deletions xarray/namedarray/dtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,19 @@

@functools.total_ordering
class AlwaysGreaterThan:
def __gt__(self, other: Any) -> Literal[True]:
def __gt__(self, other: object) -> Literal[True]:
return True

def __eq__(self, other: Any) -> bool:
def __eq__(self, other: object) -> bool:
return isinstance(other, type(self))


@functools.total_ordering
class AlwaysLessThan:
def __lt__(self, other: Any) -> Literal[True]:
def __lt__(self, other: object) -> Literal[True]:
return True

def __eq__(self, other: Any) -> bool:
def __eq__(self, other: object) -> bool:
return isinstance(other, type(self))


Expand Down
6 changes: 3 additions & 3 deletions xarray/structure/chunks.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,15 +167,15 @@ def _maybe_chunk(


@overload
def unify_chunks(__obj: _T) -> tuple[_T]: ...
def unify_chunks(obj: _T, /) -> tuple[_T]: ...


@overload
def unify_chunks(__obj1: _T, __obj2: _U) -> tuple[_T, _U]: ...
def unify_chunks(obj1: _T, obj2: _U, /) -> tuple[_T, _U]: ...


@overload
def unify_chunks(__obj1: _T, __obj2: _U, __obj3: _V) -> tuple[_T, _U, _V]: ...
def unify_chunks(obj1: _T, obj2: _U, obj3: _V, /) -> tuple[_T, _U, _V]: ...


@overload
Expand Down
5 changes: 3 additions & 2 deletions xarray/structure/merge.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
from __future__ import annotations

from collections import defaultdict
from collections.abc import Hashable, Iterable, Mapping, Sequence, Set
from collections.abc import Hashable, Iterable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from typing import TYPE_CHECKING, Any, NamedTuple, Union

import pandas as pd
Expand Down Expand Up @@ -396,7 +397,7 @@ def collect_from_coordinates(
def merge_coordinates_without_align(
objects: list[Coordinates],
prioritized: Mapping[Any, MergeElement] | None = None,
exclude_dims: Set = frozenset(),
exclude_dims: AbstractSet = frozenset(),
combine_attrs: CombineAttrsOptions = "override",
) -> tuple[dict[Hashable, Variable], dict[Hashable, Index]]:
"""Merge variables/indexes from coordinates without automatic alignments.
Expand Down
4 changes: 1 addition & 3 deletions xarray/tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -6215,9 +6215,7 @@ def test_h5netcdf_entrypoint(tmp_path: Path) -> None:

@requires_netCDF4
@pytest.mark.parametrize("str_type", (str, np.str_))
def test_write_file_from_np_str(
str_type: type[str] | type[np.str_], tmpdir: str
) -> None:
def test_write_file_from_np_str(str_type: type[str | np.str_], tmpdir: str) -> None:
# https://github.com/pydata/xarray/pull/5264
scenarios = [str_type(v) for v in ["scenario_a", "scenario_b", "scenario_c"]]
years = range(2015, 2100 + 1)
Expand Down
11 changes: 9 additions & 2 deletions xarray/util/generate_aggregations.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@

"""

import collections
import textwrap
from dataclasses import dataclass, field
from typing import NamedTuple

MODULE_PREAMBLE = '''\
"""Mixin classes with reduction operations."""
Expand Down Expand Up @@ -227,7 +227,14 @@ def {method}(
and better supported. ``cumsum`` and ``cumprod`` may be deprecated
in the future."""

ExtraKwarg = collections.namedtuple("ExtraKwarg", "docs kwarg call example")

class ExtraKwarg(NamedTuple):
docs: str
kwarg: str
call: str
example: str


skipna = ExtraKwarg(
docs=_SKIPNA_DOCSTRING,
kwarg="skipna: bool | None = None,",
Expand Down
Loading