|
6 | 6 | import pytest |
7 | 7 |
|
8 | 8 | import xarray as xr |
9 | | -from xarray import DataArray |
| 9 | +from xarray import DataArray, Dataset |
10 | 10 | from xarray.tests import assert_allclose, assert_equal |
11 | 11 |
|
12 | 12 | from . import raise_if_dask_computes, requires_cftime, requires_dask |
13 | 13 |
|
14 | 14 |
|
15 | 15 | @pytest.mark.parametrize("as_dataset", (True, False)) |
16 | | -def test_weighted_non_DataArray_weights(as_dataset): |
| 16 | +def test_weighted_non_DataArray_weights(as_dataset: bool) -> None: |
17 | 17 |
|
18 | | - data = DataArray([1, 2]) |
| 18 | + data: DataArray | Dataset = DataArray([1, 2]) |
19 | 19 | if as_dataset: |
20 | 20 | data = data.to_dataset(name="data") |
21 | 21 |
|
22 | 22 | with pytest.raises(ValueError, match=r"`weights` must be a DataArray"): |
23 | | - data.weighted([1, 2]) |
| 23 | + data.weighted([1, 2]) # type: ignore |
24 | 24 |
|
25 | 25 |
|
26 | 26 | @pytest.mark.parametrize("as_dataset", (True, False)) |
27 | 27 | @pytest.mark.parametrize("weights", ([np.nan, 2], [np.nan, np.nan])) |
28 | | -def test_weighted_weights_nan_raises(as_dataset, weights): |
| 28 | +def test_weighted_weights_nan_raises(as_dataset: bool, weights: list[float]) -> None: |
29 | 29 |
|
30 | | - data = DataArray([1, 2]) |
| 30 | + data: DataArray | Dataset = DataArray([1, 2]) |
31 | 31 | if as_dataset: |
32 | 32 | data = data.to_dataset(name="data") |
33 | 33 |
|
@@ -271,22 +271,30 @@ def test_weighted_quantile_nan(skipna): |
271 | 271 | @pytest.mark.parametrize( |
272 | 272 | "da", |
273 | 273 | ( |
274 | | - [1, 1.9, 2.2, 3, 3.7, 4.1, 5], |
275 | | - [1, 1.9, 2.2, 3, 3.7, 4.1, np.nan], |
276 | | - [np.nan, np.nan, np.nan], |
| 274 | + pytest.param([1, 1.9, 2.2, 3, 3.7, 4.1, 5], id="nonan"), |
| 275 | + pytest.param([1, 1.9, 2.2, 3, 3.7, 4.1, np.nan], id="singlenan"), |
| 276 | + pytest.param( |
| 277 | + [np.nan, np.nan, np.nan], |
| 278 | + id="allnan", |
| 279 | + marks=pytest.mark.filterwarnings( |
| 280 | + "ignore:All-NaN slice encountered:RuntimeWarning" |
| 281 | + ), |
| 282 | + ), |
277 | 283 | ), |
278 | 284 | ) |
279 | 285 | @pytest.mark.parametrize("q", (0.5, (0.2, 0.8))) |
280 | 286 | @pytest.mark.parametrize("skipna", (True, False)) |
281 | 287 | @pytest.mark.parametrize("factor", [1, 3.14]) |
282 | | -def test_weighted_quantile_equal_weights(da, q, skipna, factor): |
| 288 | +def test_weighted_quantile_equal_weights( |
| 289 | + da: list[float], q: float | tuple[float, ...], skipna: bool, factor: float |
| 290 | +) -> None: |
283 | 291 | # if all weights are equal (!= 0), should yield the same result as quantile |
284 | 292 |
|
285 | | - da = DataArray(da) |
286 | | - weights = xr.full_like(da, factor) |
| 293 | + data = DataArray(da) |
| 294 | + weights = xr.full_like(data, factor) |
287 | 295 |
|
288 | | - expected = da.quantile(q, skipna=skipna) |
289 | | - result = da.weighted(weights).quantile(q, skipna=skipna) |
| 296 | + expected = data.quantile(q, skipna=skipna) |
| 297 | + result = data.weighted(weights).quantile(q, skipna=skipna) |
290 | 298 |
|
291 | 299 | assert_allclose(expected, result) |
292 | 300 |
|
|
0 commit comments