From 044feb6e68209c3d8e4ccb2bd1f26123a55c7ebf Mon Sep 17 00:00:00 2001 From: Vo Date: Sat, 15 Aug 2026 16:03:11 -0700 Subject: [PATCH 1/2] Replace mypy with ty --- .pre-commit-config.yaml | 15 +++------------ conda-env/dev.yml | 2 +- pyproject.toml | 20 +++++++++++--------- 3 files changed, 15 insertions(+), 22 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index eebd6872..769fd33f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -25,16 +25,7 @@ repos: # Run the formatter. - id: ruff-format - - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.19.1 + - repo: https://github.com/astral-sh/ty-pre-commit + rev: v0.0.71 hooks: - - id: mypy - args: ["--config=pyproject.toml"] - additional_dependencies: - [ - dask, - "numpy>=2.0.0,<3.0.0", - pandas, - xarray>=2024.03.0, - types-python-dateutil, - ] + - id: ty diff --git a/conda-env/dev.yml b/conda-env/dev.yml index 956ecc05..1fafab3a 100644 --- a/conda-env/dev.yml +++ b/conda-env/dev.yml @@ -44,7 +44,7 @@ dependencies: # NOTE: If the tools below are updated, also update their 'rev' in `.pre-commit.config.yaml` - pre-commit=4.1.0 - ruff=0.14.10 - - mypy=1.19.1 + - ty=0.0.71 # Testing # ================== - pytest diff --git a/pyproject.toml b/pyproject.toml index fa74deef..8e358932 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,7 @@ docs = [ "ipython", "gsw-xarray", ] -dev = ["types-python-dateutil", "pre-commit", "ruff", "mypy"] +dev = ["types-python-dateutil", "pre-commit", "ruff", "ty"] [tool.setuptools.packages.find] include = ["xcdat", "xcdat.*"] @@ -125,11 +125,13 @@ python_files = ["tests.py", "test_*.py"] # These markers are defined in `xarray.tests` and must be included to avoid warnings when importing from this module. markers = ["flaky", "network"] -[tool.mypy] -# Docs: https://mypy.readthedocs.io/en/stable/config_file.html -python_version = "3.13" -check_untyped_defs = true -ignore_missing_imports = true -warn_unused_ignores = true -warn_redundant_casts = true -warn_unused_configs = true +[tool.ty.environment] +python-version = "3.13" + +[tool.ty.src] +exclude = ["docs"] + +[tool.ty.rules] +unresolved-import = "ignore" +unused-ignore-comment = "error" +redundant-cast = "error" From 81cc551ad288c1f012f56493f0b1522a6e4ed051 Mon Sep 17 00:00:00 2001 From: Vo Date: Sat, 15 Aug 2026 16:03:16 -0700 Subject: [PATCH 2/2] Fix ty diagnostics --- tests/fixtures.py | 5 +-- tests/test_axis.py | 2 +- tests/test_dataset.py | 15 ++++++++ tests/test_regrid.py | 69 ++++++++++++++++++++++++++++++++++++- tests/test_tutorial.py | 26 ++++++++++++++ xcdat/axis.py | 3 +- xcdat/bounds.py | 21 ++++++----- xcdat/dataset.py | 4 +-- xcdat/mask.py | 14 ++++---- xcdat/regridder/accessor.py | 9 +++-- xcdat/regridder/regrid2.py | 9 +++-- xcdat/regridder/xgcm.py | 60 ++++++++++++++++++++------------ xcdat/spatial.py | 2 +- xcdat/temporal.py | 35 +++++++++++-------- xcdat/tutorial.py | 17 ++++++++- 15 files changed, 224 insertions(+), 67 deletions(-) diff --git a/tests/fixtures.py b/tests/fixtures.py index bf28b807..ff026421 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -1,5 +1,6 @@ """This module stores reusable test fixtures.""" +from collections.abc import Hashable from typing import Literal import cftime @@ -457,10 +458,10 @@ def generate_multiple_variable_dataset( for idx in range(copies): ds_copy = ds_base.copy(deep=True) - var_names = list(["ts"]) + var_names: list[Hashable] = ["ts"] if separate_dims: - var_names += list(ds_base.sizes.keys()) # type: ignore[arg-type] + var_names += list(ds_base.sizes.keys()) ds_copy = ds_copy.rename({x: f"{x}{idx + 1}" for x in var_names}) diff --git a/tests/test_axis.py b/tests/test_axis.py index df604f2b..f00a7e98 100644 --- a/tests/test_axis.py +++ b/tests/test_axis.py @@ -135,7 +135,7 @@ def test_raises_error_if_dim_does_not_exist(self): for dim in dims: with pytest.raises(KeyError): - get_dim_coords(ds, dim) # type: ignore + get_dim_coords(ds, dim) def test_raises_error_if_axis_or_standard_name_is_not_set_or_dim_name_is_not_valid( self, diff --git a/tests/test_dataset.py b/tests/test_dataset.py index f5e1c912..d58648e8 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1729,3 +1729,18 @@ def test_bounds_always_persist(self): assert ds.get("lat_bnds") is not None assert ds.get("lon_bnds") is not None assert ds.get("time_bnds") is not None + + def test_drops_coords_and_dimensions_used_only_by_unselected_var(self): + ds = xr.Dataset( + data_vars={ + "selected": ("time", [1, 2]), + "unselected": ("level", [3, 4]), + }, + coords={"time": [0, 1], "level": [1000, 850]}, + ) + + result = _keep_single_var(ds, key="selected") + + assert list(result.data_vars) == ["selected"] + assert "level" not in result.coords + assert "level" not in result.dims diff --git a/tests/test_regrid.py b/tests/test_regrid.py index f78d6981..7f8d3948 100644 --- a/tests/test_regrid.py +++ b/tests/test_regrid.py @@ -1,6 +1,7 @@ import datetime import re import sys +from typing import cast from unittest import mock import numpy as np @@ -387,6 +388,50 @@ def test_missing_output_z_coord(self): ): regridder.vertical("so", ds) + def test_multiple_output_z_coords(self): + output_grid = self.output_grid.assign_coords( + ilev=xr.DataArray( + self.output_grid.lev.data, + dims="ilev", + attrs=self.output_grid.lev.attrs.copy(), + ) + ) + regridder = xgcm.XGCMRegridder( + self.ds, output_grid, method="linear", target_data=None + ) + + with pytest.raises( + RuntimeError, + match="Could not determine a single 'Z' coordinate in output dataset", + ): + regridder.vertical("so", self.ds) + + @mock.patch("xgcm.Grid") + def test_non_string_vertical_coordinate_names(self, grid): + output_coord_z = xr.DataArray( + self.output_grid.lev.data, + dims="output_lev", + name=0, + attrs=self.output_grid.lev.attrs.copy(), + ) + get_dim_coords = xgcm.get_dim_coords + + with mock.patch("xcdat.regridder.xgcm.get_dim_coords") as get_dim_coords_mock: + get_dim_coords_mock.side_effect = lambda obj, axis: ( + output_coord_z if obj is self.output_grid else get_dim_coords(obj, axis) + ) + grid.return_value.transform.return_value = self.ds.so.rename( + {"lev": "output_lev"} + ) + regridder = xgcm.XGCMRegridder( + self.ds, self.output_grid, method="linear", target_data=None + ) + + with pytest.raises( + RuntimeError, match="Vertical coordinate names must be strings" + ): + regridder.vertical("so", self.ds) + def test_missing_input_z_bounds(self): ds = fixtures.generate_lev_dataset() @@ -759,6 +804,24 @@ def test_regrid_fine_coarse_2d(self): assert np.all(output_data.ts == 1) + @pytest.mark.parametrize(("axis", "coord"), (("X", "lon"), ("Y", "lat"))) + def test_multiple_output_dimension_coords(self, axis, coord): + output_grid = self.fine_2d_ds.assign_coords( + { + f"{coord}_alternate": xr.DataArray( + self.fine_2d_ds[coord].data, + dims=f"{coord}_alternate", + attrs=self.fine_2d_ds[coord].attrs.copy(), + ) + } + ) + regridder = regrid2.Regrid2Regridder(self.coarse_2d_ds, output_grid) + + with pytest.raises( + ValueError, match=f"Multiple dimension coordinates found for {axis!r} axis" + ): + regridder.horizontal("ts", self.coarse_2d_ds) + def test_regrid_3d(self): regridder = regrid2.Regrid2Regridder(self.coarse_3d_ds, self.fine_2d_ds) @@ -1242,7 +1305,11 @@ def test_create_grid_wrong_axis_value(self): "Argument 'x' should be an xr.DataArray representing coordinates or a tuple (xr.DataArray, xr.DataArray) representing coordinates and bounds." ), ): - grid.create_grid(x=(self.lon, self.lon_bnds, self.lat)) # type: ignore[arg-type] + invalid_x = cast( + tuple[xr.DataArray, xr.DataArray | None], + (self.lon, self.lon_bnds, self.lat), + ) + grid.create_grid(x=invalid_x) def test_uniform_grid(self): new_grid = grid.create_uniform_grid(-90, 90, 4.0, -180, 180, 5.0) diff --git a/tests/test_tutorial.py b/tests/test_tutorial.py index 5e53f0ef..2055bd96 100644 --- a/tests/test_tutorial.py +++ b/tests/test_tutorial.py @@ -55,6 +55,32 @@ def test_open_dataset_no_cache(self, mock_open_dataset, mock_retrieve): ) assert not os.path.exists(mock_retrieve.return_value) + @patch("pooch.retrieve") + @patch("pooch.HTTPDownloader") + @patch("xcdat.dataset.open_dataset") + def test_retrieve_downloader_adapter( + self, mock_open_dataset, mock_http_downloader, mock_retrieve + ): + mock_open_dataset.return_value = xr.Dataset() + mock_downloader = mock_http_downloader.return_value + mock_downloader.return_value = str(self.cache_dir / "test.nc") + + def retrieve(*, downloader, **_): + return downloader( + "https://example.com/test.nc", None, None, check_only=None + ) + + mock_retrieve.side_effect = retrieve + + open_dataset("tas_amon_access", cache_dir=self.cache_dir) + + mock_downloader.assert_called_once_with( + url="https://example.com/test.nc", + output_file=None, + pooch=None, + check_only=False, + ) + def test_raises_error_with_invalid_name(self): with pytest.raises(ValueError): open_dataset("invalid_name", cache_dir=self.cache_dir) diff --git a/xcdat/axis.py b/xcdat/axis.py index 6d00a877..b794726f 100644 --- a/xcdat/axis.py +++ b/xcdat/axis.py @@ -547,7 +547,8 @@ def _align_lon_to_360( ds_lon[var.name] = new_var # Create a new dataset of non-longitude vars and updated longitude vars. - ds_no_lon = ds.get([v for v in ds.data_vars if dim not in ds[v].dims]) # type: ignore + ds_no_lon = ds.get([v for v in ds.data_vars if dim not in ds[v].dims]) + assert isinstance(ds_no_lon, xr.Dataset) ds_final = xr.merge((ds_no_lon, ds_lon)) return ds_final diff --git a/xcdat/bounds.py b/xcdat/bounds.py index 0356405b..d9319b56 100644 --- a/xcdat/bounds.py +++ b/xcdat/bounds.py @@ -3,6 +3,7 @@ import collections import datetime import warnings +from collections.abc import Callable from typing import Literal import cf_xarray as cfxr # noqa: F401 @@ -24,6 +25,10 @@ logger = _setup_custom_logger(__name__) +TimeSteps = np.ndarray | pd.DatetimeIndex +TimeBound = cftime.datetime | pd.Timestamp +TimeBoundFactory = Callable[..., TimeBound] + @xr.register_dataset_accessor("bounds") class BoundsAccessor: @@ -620,7 +625,7 @@ def _create_time_bounds( # noqa: C901 # pandas time/date components which simplifies creating bounds. # https://pandas.pydata.org/docs/user_guide/timeseries.html#time-date-components timesteps = pd.to_datetime(timesteps) - obj_type = pd.Timestamp + obj_type: TimeBoundFactory = pd.Timestamp elif _get_datetime_like_type(time) == cftime.datetime: calendar = time.encoding["calendar"] obj_type = get_date_type(calendar) @@ -666,8 +671,8 @@ def _create_time_bounds( # noqa: C901 def _create_yearly_time_bounds( self, - timesteps: np.ndarray, - obj_type: cftime.datetime | pd.Timestamp, + timesteps: TimeSteps, + obj_type: TimeBoundFactory, ) -> list[cftime.datetime | pd.Timestamp]: """Creates time bounds for each timestep with the start and end of the year. @@ -703,8 +708,8 @@ def _create_yearly_time_bounds( def _create_monthly_time_bounds( self, - timesteps: np.ndarray, - obj_type: cftime.datetime | pd.Timestamp, + timesteps: TimeSteps, + obj_type: TimeBoundFactory, end_of_month: bool = False, ) -> list[cftime.datetime | pd.Timestamp]: """Creates time bounds for each timestep with the start and end of the month. @@ -758,7 +763,7 @@ def _create_monthly_time_bounds( def _add_months_to_timestep( self, timestep: cftime.datetime | pd.Timestamp, - obj_type: cftime.datetime | pd.Timestamp, + obj_type: TimeBoundFactory, delta: int, ) -> cftime.datetime | pd.Timestamp: """Adds delta month(s) to a timestep. @@ -807,8 +812,8 @@ def _add_months_to_timestep( def _create_daily_time_bounds( self, - timesteps: np.ndarray, - obj_type: cftime.datetime | pd.Timestamp, + timesteps: TimeSteps, + obj_type: TimeBoundFactory, freq: Literal[1, 2, 3, 4, 6, 8, 12, 24] = 1, ) -> list[cftime.datetime | pd.Timestamp]: """Creates time bounds for each timestep with the start and end of the day. diff --git a/xcdat/dataset.py b/xcdat/dataset.py index 742a188e..8ab78354 100644 --- a/xcdat/dataset.py +++ b/xcdat/dataset.py @@ -6,7 +6,7 @@ from datetime import datetime from functools import partial from io import BufferedIOBase -from typing import Any, Literal +from typing import Any, Literal, cast import numpy as np import xarray as xr @@ -769,7 +769,7 @@ def _keep_single_var(dataset: xr.Dataset, key: str) -> xr.Dataset: if key in bounds_vars: raise ValueError("Please specify a non-bounds data variable.") - return dataset[[key] + bounds_vars] + return cast(xr.Dataset, dataset[[key] + bounds_vars]) def _get_data_var(dataset: xr.Dataset, key: str) -> xr.DataArray: diff --git a/xcdat/mask.py b/xcdat/mask.py index bcbead9b..28d09bf7 100644 --- a/xcdat/mask.py +++ b/xcdat/mask.py @@ -167,7 +167,8 @@ def generate_land_sea_mask( if method == "regionmask": land_mask = regionmask.defined_regions.natural_earth_v5_0_0.land_110 - lon, lat = get_dim_coords(da, "X"), get_dim_coords(da, "Y") + lon = _as_dataarray(get_dim_coords(da, "X")) + lat = _as_dataarray(get_dim_coords(da, "Y")) land_sea_mask = land_mask.mask(lon, lat=lat) @@ -261,11 +262,6 @@ def pcmdi_land_sea_mask( >>> from xcdat._data import _get_pcmdi_mask_path >>> path = _get_pcmdi_mask_path() """ - if source is not None and source_data_var is None: - raise ValueError( - "The 'source_data_var' value cannot be None when using the 'source' option." - ) - if source is None: source_data_var = "sftlf" @@ -274,6 +270,10 @@ def pcmdi_land_sea_mask( # Turn off time decoding to prevent logger warning since this dataset # does not have a time axis. source = open_dataset(resource_path, decode_times=False) + elif source_data_var is None: + raise ValueError( + "The 'source_data_var' value cannot be None when using the 'source' option." + ) source_regrid = source.regridder.horizontal( source_data_var, _obj_to_grid_ds(da), tool="regrid2" @@ -302,7 +302,7 @@ def pcmdi_land_sea_mask( improved_mask = _improve_mask( mask.copy(deep=True), source_regrid, - source_data_var, # type: ignore[arg-type] + source_data_var, surrounds, is_circular, threshold1, diff --git a/xcdat/regridder/accessor.py b/xcdat/regridder/accessor.py index 132d8d9d..8f094114 100644 --- a/xcdat/regridder/accessor.py +++ b/xcdat/regridder/accessor.py @@ -333,8 +333,11 @@ def _get_axis_coord_and_bounds( ) except (ValueError, KeyError): try: - coord_var = get_dim_coords(obj, axis, multidim=multidim) # type: ignore - _validate_grid_has_single_axis_dim(axis, coord_var) + coord_var = get_dim_coords(obj, axis, multidim=multidim) + if isinstance(coord_var, xr.Dataset): + _validate_grid_has_single_axis_dim(axis, coord_var) + + assert isinstance(coord_var, xr.DataArray) except KeyError: coord_var = None @@ -343,7 +346,7 @@ def _get_axis_coord_and_bounds( bounds_var = None bounds_key = coord_var.attrs.get("bounds") - if bounds_key: + if isinstance(bounds_key, str) and bounds_key: try: bounds_var = obj.get(bounds_key) except AttributeError: diff --git a/xcdat/regridder/regrid2.py b/xcdat/regridder/regrid2.py index dfa90eab..aa97007c 100644 --- a/xcdat/regridder/regrid2.py +++ b/xcdat/regridder/regrid2.py @@ -392,10 +392,13 @@ def _get_output_coords( # First get the X and Y axes from the output grid. for key in ["X", "Y"]: - input_coord = xc.get_dim_coords(dv_input, key) # type: ignore - output_coord = xc.get_dim_coords(output_grid, key) # type: ignore + input_coord = xc.get_dim_coords(dv_input, key) + output_coord = xc.get_dim_coords(output_grid, key) - output_coords[str(input_coord.name)] = output_coord # type: ignore + if isinstance(output_coord, xr.Dataset): + raise ValueError(f"Multiple dimension coordinates found for {key!r} axis") + + output_coords[str(input_coord.name)] = output_coord # Get the remaining axes the input data variable (e.g., "time"). for dim in dv_input.dims: diff --git a/xcdat/regridder/xgcm.py b/xcdat/regridder/xgcm.py index 9b3552e2..ce22e568 100644 --- a/xcdat/regridder/xgcm.py +++ b/xcdat/regridder/xgcm.py @@ -1,4 +1,3 @@ -from collections.abc import Hashable from typing import Any, Literal, get_args import xarray as xr @@ -175,6 +174,12 @@ def vertical(self, data_var: str, ds: xr.Dataset) -> xr.Dataset: "Could not determine 'Z' coordinate in output dataset" ) from e + if isinstance(output_coord_z, xr.Dataset): + raise RuntimeError( + "Could not determine a single 'Z' coordinate in output dataset" + ) + + grid_coords: dict[str, dict[str, str]] if self._grid_positions is None: grid_coords = self._get_grid_positions() else: @@ -208,10 +213,19 @@ def vertical(self, data_var: str, ds: xr.Dataset) -> xr.Dataset: # transposed to match the input dimension order if output_da.dims != ds[data_var].dims: input_coord_z = get_dim_coords(ds[data_var], "Z") + input_coord_z_name = input_coord_z.name + output_coord_z_name = output_coord_z.name + + if not isinstance(input_coord_z_name, str) or not isinstance( + output_coord_z_name, str + ): + raise RuntimeError("Vertical coordinate names must be strings") output_order = [ - x.replace(input_coord_z.name, output_coord_z.name) # type: ignore[attr-defined] - for x in ds[data_var].dims + dim.replace(input_coord_z_name, output_coord_z_name) + if isinstance(dim, str) + else dim + for dim in ds[data_var].dims ] output_da = output_da.transpose(*output_order) @@ -284,14 +298,13 @@ def _infer_target_data(self, ds) -> xr.DataArray | None: return ds.decoded_vertical_coord - def _get_target_data(self, ds) -> xr.DataArray | None: + def _get_target_data(self, ds: xr.Dataset) -> xr.DataArray | None: """Retrieve the target data from the given xarray Dataset. Attempts to access the target data variable from the provided dataset using the attribute `self._target_data`. If `self._target_data` is a - string and not found in the dataset, raises a RuntimeError. If - `self._target_data` is not a string or is None, returns None. If a - ValueError occurs, returns `self._target_data` as is. + string and not found in the dataset, raises a RuntimeError. If it is a + DataArray, returns it as is. If it is None, returns None. Parameters ---------- @@ -308,21 +321,20 @@ def _get_target_data(self, ds) -> xr.DataArray | None: RuntimeError If `self._target_data` is a string and not found in the dataset. """ - try: - target_data = ds[self._target_data] - except ValueError: - target_data = self._target_data - except KeyError as e: - if self._target_data is not None and isinstance(self._target_data, str): - raise RuntimeError( - f"Could not find target variable {self._target_data!r} in dataset" - ) from e + if self._target_data is None: + return None - target_data = None + if isinstance(self._target_data, xr.DataArray): + return self._target_data - return target_data + try: + return ds[self._target_data] + except KeyError as e: + raise RuntimeError( + f"Could not find target variable {self._target_data!r} in dataset" + ) from e - def _get_grid_positions(self) -> dict[str, Any | Hashable]: + def _get_grid_positions(self) -> dict[str, dict[str, str]]: """ Determine the grid point positions for the "Z" axis in the input grid. @@ -334,7 +346,7 @@ def _get_grid_positions(self) -> dict[str, Any | Hashable]: Returns ------- - dict[str, Any | Hashable] + dict[str, dict[str, str]] Mapping of the "Z" axis to its grid position, e.g., {"Z": {"center": }}. @@ -357,13 +369,14 @@ def _get_grid_positions(self) -> dict[str, Any | Hashable]: raise RuntimeError("Could not determine `Z` coordinate in dataset.") from e if isinstance(coord_z, xr.Dataset): - coords = ", ".join(sorted(list(coord_z.coords.keys()))) # type: ignore[arg-type] + coord_names = list(coord_z.coords.keys()) + coords = ", ".join(sorted(str(coord) for coord in coord_names)) raise RuntimeError( "Could not determine the `Z` coordinate in the input grid. " f"Found multiple axes ({coords}), ensure there is only a " "single `Z` axis in the input grid.", - list(coord_z.coords.keys()), + coord_names, ) try: @@ -371,6 +384,9 @@ def _get_grid_positions(self) -> dict[str, Any | Hashable]: except KeyError as e: raise RuntimeError("Could not determine `Z` bounds in dataset.") from e + if not isinstance(coord_z.name, str): + raise RuntimeError("Vertical coordinate name must be a string") + # handle simple point positions based on point and bounds if (coord_z[0] > bounds_z[0][0] and coord_z[0] < bounds_z[0][1]) or ( coord_z[0] < bounds_z[0][0] and coord_z[0] > bounds_z[0][1] diff --git a/xcdat/spatial.py b/xcdat/spatial.py index def74c99..65d8dd84 100644 --- a/xcdat/spatial.py +++ b/xcdat/spatial.py @@ -872,7 +872,7 @@ def _combine_weights(self, axis_weights: AxisWeights) -> xr.DataArray: region_weights = reduce((lambda x, y: x * y), axis_weights.values()) coord_keys = sorted(region_weights.dims) # type: ignore - region_weights.name = "_".join(coord_keys) + "_wts" # type: ignore + region_weights.name = "_".join(coord_keys) + "_wts" return region_weights diff --git a/xcdat/temporal.py b/xcdat/temporal.py index a03a933d..8d634236 100644 --- a/xcdat/temporal.py +++ b/xcdat/temporal.py @@ -11,6 +11,7 @@ import pandas as pd import xarray as xr from dask.array.core import Array +from pandas._libs.tslibs.nattype import NaTType from xarray.coding.cftime_offsets import get_date_type from xarray.core.common import contains_cftime_datetimes, is_np_datetime_like from xarray.core.groupby import DataArrayGroupBy @@ -1195,16 +1196,14 @@ def _preprocess_dataset(self, ds: xr.Dataset) -> xr.Dataset: {self.dim: slice(self._reference_period[0], self._reference_period[1])} ) - if ( - self._freq == "season" - and self._season_config.get("custom_seasons") is not None - ): + custom_seasons = self._season_config.get("custom_seasons") + if self._freq == "season" and custom_seasons is not None: # Get a flat list of all of the months included in the custom # seasons to determine if the dataset needs to be subsetted # on just those months. For example, if we define a custom season # "NDJFM", we should subset the dataset for time coordinates # belonging to those months. - months = self._season_config["custom_seasons"].values() # type: ignore + months = custom_seasons.values() months = list(chain.from_iterable(months)) if len(months) != 12: @@ -1295,12 +1294,13 @@ def _shift_custom_season_years(self, ds: xr.Dataset) -> xr.Dataset: """ ds_new = ds.copy() custom_seasons = self._season_config["custom_seasons"] + assert custom_seasons is not None # Identify months that span across years in custom seasons by getting # the months before "Jan" if "Jan" is not the first month of the season. # Note: Only one custom season can span the calendar year. span_months: list[int] = [] - for months in custom_seasons.values(): # type: ignore + for months in custom_seasons.values(): month_ints = [MONTH_STR_TO_INT[month] for month in months] if 1 in month_ints and month_ints.index(1) != 0: @@ -1404,7 +1404,7 @@ def _shift_cftime_year(self, time: cftime.datetime) -> cftime.datetime: """ return time.replace(year=time.year + 1) - def _shift_datetime_year(self, time) -> pd.Timestamp: + def _shift_datetime_year(self, time) -> pd.Timestamp | NaTType: """ Shift the year of a datetime-like object by 1. @@ -1415,7 +1415,7 @@ def _shift_datetime_year(self, time) -> pd.Timestamp: Returns ------- - pd.Timestamp + pd.Timestamp | NaTType The datetime-like object with the year incremented by 1. """ ts = pd.Timestamp(time) @@ -1528,13 +1528,17 @@ def _drop_incomplete_seasons(self, ds: xr.Dataset) -> xr.Dataset: # broadcasting, which is a behavior we do not desire. # https://github.com/pydata/xarray/issues/1234 # https://github.com/pydata/xarray/issues/8796#issuecomment-1974878267 - ds_no_time = ds.get([v for v in ds.data_vars if self.dim not in ds[v].dims]) # type: ignore - ds_time = ds.get([v for v in ds.data_vars if self.dim in ds[v].dims]) # type: ignore + ds_no_time = ds.get([v for v in ds.data_vars if self.dim not in ds[v].dims]) + ds_time = ds.get([v for v in ds.data_vars if self.dim in ds[v].dims]) + assert isinstance(ds_no_time, xr.Dataset) + assert isinstance(ds_time, xr.Dataset) coords_to_drop = time_coords.values[indexes_to_drop] ds_time = ds_time.where(~time_coords.isin(coords_to_drop), drop=True) + assert isinstance(ds_time, xr.Dataset) - ds_new = xr.merge([ds_time, ds_no_time]) + ds_new = xr.merge((ds_time, ds_no_time)) + assert isinstance(ds_new, xr.Dataset) return ds_new @@ -1976,12 +1980,13 @@ def _map_months_to_custom_seasons(self, df: pd.DataFrame) -> pd.DataFrame: to a custom season. """ custom_seasons = self._season_config["custom_seasons"] + assert custom_seasons is not None # NOTE: This for loop has a time complexity of O(n^2), but it is fine # because these data structures are small. seasons_map = {} for mon_int, mon_str in MONTH_INT_TO_STR.items(): - for season in custom_seasons: # type: ignore + for season in custom_seasons: if mon_str in season: seasons_map[mon_int] = season @@ -2164,7 +2169,7 @@ def _add_operation_attrs(self, data_var: xr.DataArray) -> xr.DataArray: } if self._weighted: - attrs_to_set["min_weight"] = self._min_weight # type: ignore + attrs_to_set["min_weight"] = self._min_weight if self._freq == "season": drop_incomplete_seasons = self._season_config["drop_incomplete_seasons"] @@ -2179,10 +2184,10 @@ def _add_operation_attrs(self, data_var: xr.DataArray) -> xr.DataArray: custom_seasons = self._season_config.get("custom_seasons") if custom_seasons is not None: - attrs_to_set["custom_seasons"] = list(custom_seasons.keys()) # type: ignore + attrs_to_set["custom_seasons"] = list(custom_seasons.keys()) else: dec_mode = self._season_config.get("dec_mode") - attrs_to_set["dec_mode"] = dec_mode # type: ignore + attrs_to_set["dec_mode"] = dec_mode data_var.attrs.update(attrs_to_set) diff --git a/xcdat/tutorial.py b/xcdat/tutorial.py index 463122e7..bad0fca7 100644 --- a/xcdat/tutorial.py +++ b/xcdat/tutorial.py @@ -8,6 +8,7 @@ import os import pathlib import sys +from typing import Any import xarray as xr from xarray.tutorial import _construct_cache_dir, file_formats @@ -107,8 +108,22 @@ def open_dataset( headers = {"User-Agent": f"xcdat {sys.modules['xcdat'].__version__}"} downloader = pooch.HTTPDownloader(headers=headers) + def download( + fname: str, + action: str | os.PathLike[str] | None, + pooch: Any, + *, + check_only: bool | None = None, + ) -> Any: + return downloader( + url=fname, + output_file=action, + pooch=pooch, + check_only=bool(check_only), + ) + filepath = pooch.retrieve( - url=url, known_hash=None, path=cache_dir, downloader=downloader + url=url, known_hash=None, path=cache_dir, downloader=download ) ds = open_dataset(filepath, **kargs, add_bounds=add_bounds)