diff --git a/.github/workflows/ci-additional.yaml b/.github/workflows/ci-additional.yaml index 9aa3b17746f..bb39f1875e9 100644 --- a/.github/workflows/ci-additional.yaml +++ b/.github/workflows/ci-additional.yaml @@ -35,14 +35,13 @@ jobs: runs-on: "ubuntu-latest" needs: detect-ci-trigger if: needs.detect-ci-trigger.outputs.triggered == 'false' + defaults: run: shell: bash -l {0} - env: CONDA_ENV_FILE: ci/requirements/environment.yml PYTHON_VERSION: "3.11" - steps: - uses: actions/checkout@v4 with: @@ -128,7 +127,7 @@ jobs: python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report xarray/ - name: Upload mypy coverage to Codecov - uses: codecov/codecov-action@v4.1.0 + uses: codecov/codecov-action@v4.2.0 with: file: mypy_report/cobertura.xml flags: mypy @@ -182,7 +181,7 @@ jobs: python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report xarray/ - name: Upload mypy coverage to Codecov - uses: codecov/codecov-action@v4.1.0 + uses: codecov/codecov-action@v4.2.0 with: file: mypy_report/cobertura.xml flags: mypy39 @@ -243,7 +242,7 @@ jobs: python -m pyright xarray/ - name: Upload pyright coverage to Codecov - uses: codecov/codecov-action@v4.1.0 + uses: codecov/codecov-action@v4.2.0 with: file: pyright_report/cobertura.xml flags: pyright @@ -302,7 +301,7 @@ jobs: python -m pyright xarray/ - name: Upload pyright coverage to Codecov - uses: codecov/codecov-action@v4.1.0 + uses: codecov/codecov-action@v4.2.0 with: file: pyright_report/cobertura.xml flags: pyright39 diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index a37ff876e20..9724c18436e 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -127,6 +127,14 @@ jobs: run: | python -c "import xarray" + - name: Restore cached hypothesis directory + uses: actions/cache@v4 + with: + path: .hypothesis/ + key: cache-hypothesis + enableCrossOsArchive: true + save-always: true + - name: Run tests run: python -m pytest -n 4 --timeout 180 @@ -143,7 +151,7 @@ jobs: path: pytest.xml - name: Upload code coverage to Codecov - uses: codecov/codecov-action@v4.1.0 + uses: codecov/codecov-action@v4.2.0 with: file: ./coverage.xml flags: unittests diff --git a/.github/workflows/hypothesis.yaml b/.github/workflows/hypothesis.yaml new file mode 100644 index 00000000000..2772dac22b1 --- /dev/null +++ b/.github/workflows/hypothesis.yaml @@ -0,0 +1,116 @@ +name: Slow Hypothesis CI +on: + push: + branches: + - "main" + pull_request: + branches: + - "main" + types: [opened, reopened, synchronize, labeled] + schedule: + - cron: "0 0 * * *" # Daily “At 00:00” UTC + workflow_dispatch: # allows you to trigger manually + +jobs: + detect-ci-trigger: + name: detect ci trigger + runs-on: ubuntu-latest + if: | + github.repository == 'pydata/xarray' + && (github.event_name == 'push' || github.event_name == 'pull_request') + outputs: + triggered: ${{ steps.detect-trigger.outputs.trigger-found }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 2 + - uses: xarray-contrib/ci-trigger@v1 + id: detect-trigger + with: + keyword: "[skip-ci]" + + hypothesis: + name: Slow Hypothesis Tests + runs-on: "ubuntu-latest" + needs: detect-ci-trigger + if: | + always() + && ( + (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') + || needs.detect-ci-trigger.outputs.triggered == 'true' + || contains( github.event.pull_request.labels.*.name, 'run-slow-hypothesis') + ) + defaults: + run: + shell: bash -l {0} + + env: + CONDA_ENV_FILE: ci/requirements/environment.yml + PYTHON_VERSION: "3.12" + + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 # Fetch all history for all branches and tags. + + - name: set environment variables + run: | + echo "TODAY=$(date +'%Y-%m-%d')" >> $GITHUB_ENV + + - name: Setup micromamba + uses: mamba-org/setup-micromamba@v1 + with: + environment-file: ci/requirements/environment.yml + environment-name: xarray-tests + create-args: >- + python=${{env.PYTHON_VERSION}} + pytest-reportlog + cache-environment: true + cache-environment-key: "${{runner.os}}-${{runner.arch}}-py${{env.PYTHON_VERSION}}-${{env.TODAY}}-${{hashFiles(env.CONDA_ENV_FILE)}}" + + - name: Install xarray + run: | + python -m pip install --no-deps -e . + - name: Version info + run: | + conda info -a + conda list + python xarray/util/print_versions.py + + # https://github.com/actions/cache/blob/main/tips-and-workarounds.md#update-a-cache + - name: Restore cached hypothesis directory + id: restore-hypothesis-cache + uses: actions/cache/restore@v4 + with: + path: .hypothesis/ + key: cache-hypothesis-${{ runner.os }}-${{ github.run_id }} + restore-keys: | + cache-hypothesis- + + - name: Run slow Hypothesis tests + if: success() + id: status + run: | + python -m pytest --hypothesis-show-statistics --run-slow-hypothesis properties/*.py \ + --report-log output-${{ matrix.python-version }}-log.jsonl + + # explicitly save the cache so it gets updated, also do this even if it fails. + - name: Save cached hypothesis directory + id: save-hypothesis-cache + if: always() && steps.status.outcome != 'skipped' + uses: actions/cache/save@v4 + with: + path: .hypothesis/ + key: cache-hypothesis-${{ runner.os }}-${{ github.run_id }} + + - name: Generate and publish the report + if: | + failure() + && steps.status.outcome == 'failure' + && github.event_name == 'schedule' + && github.repository_owner == 'pydata' + uses: xarray-contrib/issue-from-pytest-log@v1 + with: + log-path: output-${{ matrix.python-version }}-log.jsonl + issue-title: "Nightly Hypothesis tests failed" + issue-label: "topic-hypothesis" diff --git a/.github/workflows/upstream-dev-ci.yaml b/.github/workflows/upstream-dev-ci.yaml index 872b2d865fb..9d43c575b1a 100644 --- a/.github/workflows/upstream-dev-ci.yaml +++ b/.github/workflows/upstream-dev-ci.yaml @@ -143,7 +143,7 @@ jobs: run: | python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report - name: Upload mypy coverage to Codecov - uses: codecov/codecov-action@v4.1.0 + uses: codecov/codecov-action@v4.2.0 with: file: mypy_report/cobertura.xml flags: mypy diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 74d77e2f2ca..970b2e5e8ca 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -13,13 +13,13 @@ repos: - id: mixed-line-ending - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: 'v0.2.0' + rev: 'v0.3.4' hooks: - id: ruff args: ["--fix", "--show-fixes"] # https://github.com/python/black#version-control-integration - repo: https://github.com/psf/black-pre-commit-mirror - rev: 24.1.1 + rev: 24.3.0 hooks: - id: black-jupyter - repo: https://github.com/keewis/blackdoc @@ -27,10 +27,10 @@ repos: hooks: - id: blackdoc exclude: "generate_aggregations.py" - additional_dependencies: ["black==24.1.1"] + additional_dependencies: ["black==24.3.0"] - id: blackdoc-autoupdate-black - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.8.0 + rev: v1.9.0 hooks: - id: mypy # Copied from setup.cfg diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 4b557935dda..f62954a4176 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -15,14 +15,44 @@ What's New np.random.seed(123456) -.. _whats-new.2024.03.0: +.. _whats-new.2024.04.0: -v2024.03.0 (unreleased) +v2024.04.0 (unreleased) ----------------------- New Features ~~~~~~~~~~~~ + +Breaking changes +~~~~~~~~~~~~~~~~ + + +Bug fixes +~~~~~~~~~ + + +Internal Changes +~~~~~~~~~~~~~~~~ + + +.. _whats-new.2024.03.0: + +v2024.03.0 (Mar 29, 2024) +------------------------- + +This release brings performance improvements for grouped and resampled quantile calculations, CF decoding improvements, +minor optimizations to distributed Zarr writes, and compatibility fixes for Numpy 2.0 and Pandas 3.0. + +Thanks to the 18 contributors to this release: +Anderson Banihirwe, Christoph Hasse, Deepak Cherian, Etienne Schalk, Justus Magin, Kai Mühlbauer, Kevin Schwarzwald, Mark Harfouche, Martin, Matt Savoie, Maximilian Roos, Ray Bell, Roberto Chang, Spencer Clark, Tom Nicholas, crusaderky, owenlittlejohns, saschahofmann + +New Features +~~~~~~~~~~~~ +- Partial writes to existing chunks with ``region`` or ``append_dim`` will now raise an error + (unless ``safe_chunks=False``); previously an error would only be raised on + new variables. (:pull:`8459`, :issue:`8371`, :issue:`8882`) + By `Maximilian Roos `_. - Grouped and resampling quantile calculations now use the vectorized algorithm in ``flox>=0.9.4`` if present. By `Deepak Cherian `_. - Do not broadcast in arithmetic operations when global option ``arithmetic_broadcast=False`` @@ -37,16 +67,14 @@ New Features - Allow creating :py:class:`xr.Coordinates` objects with no indexes (:pull:`8711`) By `Benoit Bovy `_ and `Tom Nicholas `_. +- Enable plotting of ``datetime.dates``. (:issue:`8866`, :pull:`8873`) + By `Sascha Hofmann `_. Breaking changes ~~~~~~~~~~~~~~~~ - - Don't allow overwriting index variables with ``to_zarr`` region writes. (:issue:`8589`, :pull:`8876`). By `Deepak Cherian `_. -Deprecations -~~~~~~~~~~~~ - Bug fixes ~~~~~~~~~ @@ -73,11 +101,9 @@ Bug fixes - Warn and return bytes undecoded in case of UnicodeDecodeError in h5netcdf-backend (:issue:`5563`, :pull:`8874`). By `Kai Mühlbauer `_. - - -Documentation -~~~~~~~~~~~~~ - +- Fix bug incorrectly disallowing creation of a dataset with a multidimensional coordinate variable with the same name as one of its dims. + (:issue:`8884`, :pull:`8886`) + By `Tom Nicholas `_. Internal Changes ~~~~~~~~~~~~~~~~ diff --git a/properties/conftest.py b/properties/conftest.py index 0a66d92ebc6..30e638161a1 100644 --- a/properties/conftest.py +++ b/properties/conftest.py @@ -1,3 +1,24 @@ +import pytest + + +def pytest_addoption(parser): + parser.addoption( + "--run-slow-hypothesis", + action="store_true", + default=False, + help="run slow hypothesis tests", + ) + + +def pytest_collection_modifyitems(config, items): + if config.getoption("--run-slow-hypothesis"): + return + skip_slow_hyp = pytest.mark.skip(reason="need --run-slow-hypothesis option to run") + for item in items: + if "slow_hypothesis" in item.keywords: + item.add_marker(skip_slow_hyp) + + try: from hypothesis import settings except ImportError: diff --git a/properties/test_index_manipulation.py b/properties/test_index_manipulation.py new file mode 100644 index 00000000000..77b7fcbcd99 --- /dev/null +++ b/properties/test_index_manipulation.py @@ -0,0 +1,273 @@ +import itertools + +import numpy as np +import pytest + +import xarray as xr +from xarray import Dataset +from xarray.testing import _assert_internal_invariants + +pytest.importorskip("hypothesis") +pytestmark = pytest.mark.slow_hypothesis + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st +from hypothesis import note, settings +from hypothesis.stateful import ( + RuleBasedStateMachine, + initialize, + invariant, + precondition, + rule, +) + +import xarray.testing.strategies as xrst + + +@st.composite +def unique(draw, strategy): + # https://stackoverflow.com/questions/73737073/create-hypothesis-strategy-that-returns-unique-values + seen = draw(st.shared(st.builds(set), key="key-for-unique-elems")) + return draw( + strategy.filter(lambda x: x not in seen).map(lambda x: seen.add(x) or x) + ) + + +# Share to ensure we get unique names on each draw, +# so we don't try to add two variables with the same name +# or stack to a dimension with a name that already exists in the Dataset. +UNIQUE_NAME = unique(strategy=xrst.names()) +DIM_NAME = xrst.dimension_names(name_strategy=UNIQUE_NAME, min_dims=1, max_dims=1) +index_variables = st.builds( + xr.Variable, + data=npst.arrays( + dtype=xrst.pandas_index_dtypes(), + shape=npst.array_shapes(min_dims=1, max_dims=1), + elements=dict(allow_nan=False, allow_infinity=False, allow_subnormal=False), + unique=True, + ), + dims=DIM_NAME, + attrs=xrst.attrs(), +) + + +def add_dim_coord_and_data_var(ds, var): + (name,) = var.dims + # dim coord + ds[name] = var + # non-dim coord of same size; this allows renaming + ds[name + "_"] = var + + +class DatasetStateMachine(RuleBasedStateMachine): + # Can't use bundles because we'd need pre-conditions on consumes(bundle) + # indexed_dims = Bundle("indexed_dims") + # multi_indexed_dims = Bundle("multi_indexed_dims") + + def __init__(self): + super().__init__() + self.dataset = Dataset() + self.check_default_indexes = True + + # We track these separately as lists so we can guarantee order of iteration over them. + # Order of iteration over Dataset.dims is not guaranteed + self.indexed_dims = [] + self.multi_indexed_dims = [] + + @initialize(var=index_variables) + def init_ds(self, var): + """Initialize the Dataset so that at least one rule will always fire.""" + (name,) = var.dims + add_dim_coord_and_data_var(self.dataset, var) + + self.indexed_dims.append(name) + + # TODO: stacking with a timedelta64 index and unstacking converts it to object + @rule(var=index_variables) + def add_dim_coord(self, var): + (name,) = var.dims + note(f"adding dimension coordinate {name}") + add_dim_coord_and_data_var(self.dataset, var) + + self.indexed_dims.append(name) + + @rule(var=index_variables) + def assign_coords(self, var): + (name,) = var.dims + note(f"assign_coords: {name}") + self.dataset = self.dataset.assign_coords({name: var}) + + self.indexed_dims.append(name) + + @property + def has_indexed_dims(self) -> bool: + return bool(self.indexed_dims + self.multi_indexed_dims) + + @rule(data=st.data()) + @precondition(lambda self: self.has_indexed_dims) + def reset_index(self, data): + dim = data.draw(st.sampled_from(self.indexed_dims + self.multi_indexed_dims)) + self.check_default_indexes = False + note(f"> resetting {dim}") + self.dataset = self.dataset.reset_index(dim) + + if dim in self.indexed_dims: + del self.indexed_dims[self.indexed_dims.index(dim)] + elif dim in self.multi_indexed_dims: + del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)] + + @rule(newname=UNIQUE_NAME, data=st.data(), create_index=st.booleans()) + @precondition(lambda self: bool(self.indexed_dims)) + def stack(self, newname, data, create_index): + oldnames = data.draw( + st.lists( + st.sampled_from(self.indexed_dims), + min_size=1, + max_size=3 if create_index else None, + unique=True, + ) + ) + note(f"> stacking {oldnames} as {newname}") + self.dataset = self.dataset.stack( + {newname: oldnames}, create_index=create_index + ) + + if create_index: + self.multi_indexed_dims += [newname] + + # if create_index is False, then we just drop these + for dim in oldnames: + del self.indexed_dims[self.indexed_dims.index(dim)] + + @rule(data=st.data()) + @precondition(lambda self: bool(self.multi_indexed_dims)) + def unstack(self, data): + # TODO: add None + dim = data.draw(st.sampled_from(self.multi_indexed_dims)) + note(f"> unstacking {dim}") + if dim is not None: + pd_index = self.dataset.xindexes[dim].index + self.dataset = self.dataset.unstack(dim) + + del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)] + + if dim is not None: + self.indexed_dims.extend(pd_index.names) + else: + # TODO: fix this + pass + + @rule(newname=UNIQUE_NAME, data=st.data()) + @precondition(lambda self: bool(self.dataset.variables)) + def rename_vars(self, newname, data): + dim = data.draw(st.sampled_from(sorted(self.dataset.variables))) + # benbovy: "skip the default indexes invariant test when the name of an + # existing dimension coordinate is passed as input kwarg or dict key + # to .rename_vars()." + self.check_default_indexes = False + note(f"> renaming {dim} to {newname}") + self.dataset = self.dataset.rename_vars({dim: newname}) + + if dim in self.indexed_dims: + del self.indexed_dims[self.indexed_dims.index(dim)] + elif dim in self.multi_indexed_dims: + del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)] + + @precondition(lambda self: bool(self.dataset.dims)) + @rule(data=st.data()) + def drop_dims(self, data): + dims = data.draw( + st.lists( + st.sampled_from(sorted(tuple(self.dataset.dims))), + min_size=1, + unique=True, + ) + ) + note(f"> drop_dims: {dims}") + self.dataset = self.dataset.drop_dims(dims) + + for dim in dims: + if dim in self.indexed_dims: + del self.indexed_dims[self.indexed_dims.index(dim)] + elif dim in self.multi_indexed_dims: + del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)] + + @precondition(lambda self: bool(self.indexed_dims)) + @rule(data=st.data()) + def drop_indexes(self, data): + self.check_default_indexes = False + + dims = data.draw( + st.lists(st.sampled_from(self.indexed_dims), min_size=1, unique=True) + ) + note(f"> drop_indexes: {dims}") + self.dataset = self.dataset.drop_indexes(dims) + + for dim in dims: + if dim in self.indexed_dims: + del self.indexed_dims[self.indexed_dims.index(dim)] + elif dim in self.multi_indexed_dims: + del self.multi_indexed_dims[self.multi_indexed_dims.index(dim)] + + @property + def swappable_dims(self): + ds = self.dataset + options = [] + for dim in self.indexed_dims: + choices = [ + name + for name, var in ds._variables.items() + if var.dims == (dim,) + # TODO: Avoid swapping a dimension to itself + and name != dim + ] + options.extend( + (a, b) for a, b in itertools.zip_longest((dim,), choices, fillvalue=dim) + ) + return options + + @rule(data=st.data()) + # TODO: swap_dims is basically all broken if a multiindex is present + # TODO: Avoid swapping from Index to a MultiIndex level + # TODO: Avoid swapping from MultiIndex to a level of the same MultiIndex + # TODO: Avoid swapping when a MultiIndex is present + @precondition(lambda self: not bool(self.multi_indexed_dims)) + @precondition(lambda self: bool(self.swappable_dims)) + def swap_dims(self, data): + ds = self.dataset + options = self.swappable_dims + dim, to = data.draw(st.sampled_from(options)) + note( + f"> swapping {dim} to {to}, found swappable dims: {options}, all_dims: {tuple(self.dataset.dims)}" + ) + self.dataset = ds.swap_dims({dim: to}) + + del self.indexed_dims[self.indexed_dims.index(dim)] + self.indexed_dims += [to] + + @invariant() + def assert_invariants(self): + # note(f"> ===\n\n {self.dataset!r} \n===\n\n") + _assert_internal_invariants(self.dataset, self.check_default_indexes) + + +DatasetStateMachine.TestCase.settings = settings(max_examples=300, deadline=None) +DatasetTest = DatasetStateMachine.TestCase + + +@pytest.mark.skip(reason="failure detected by hypothesis") +def test_unstack_object(): + import xarray as xr + + ds = xr.Dataset() + ds["0"] = np.array(["", "\x000"], dtype=object) + ds.stack({"1": ["0"]}).unstack() + + +@pytest.mark.skip(reason="failure detected by hypothesis") +def test_unstack_timedelta_index(): + import xarray as xr + + ds = xr.Dataset() + ds["0"] = np.array([0, 1, 2, 3], dtype="timedelta64[ns]") + ds.stack({"1": ["0"]}).unstack() diff --git a/pyproject.toml b/pyproject.toml index 23595620a79..fe300cb494f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -172,12 +172,10 @@ module = [ "xarray.tests.test_dask", "xarray.tests.test_dataarray", "xarray.tests.test_duck_array_ops", - "xarray.tests.test_groupby", "xarray.tests.test_indexing", "xarray.tests.test_merge", "xarray.tests.test_missing", "xarray.tests.test_parallelcompat", - "xarray.tests.test_plot", "xarray.tests.test_sparse", "xarray.tests.test_ufuncs", "xarray.tests.test_units", @@ -296,6 +294,7 @@ markers = [ "flaky: flaky tests", "network: tests requiring a network connection", "slow: slow tests", + "slow_hypothesis: slow hypothesis tests", ] minversion = "7" python_files = "test_*.py" diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 13b1819f206..3d6baeefe01 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -195,7 +195,7 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name, safe_chunks): f"Writing this array in parallel with dask could lead to corrupted data." ) if safe_chunks: - raise NotImplementedError( + raise ValueError( base_error + " Consider either rechunking using `chunk()`, deleting " "or modifying `encoding['chunks']`, or specify `safe_chunks=False`." @@ -623,7 +623,12 @@ def store( # avoid needing to load index variables into memory. # TODO: consider making loading indexes lazy again? existing_vars, _, _ = conventions.decode_cf_variables( - self.get_variables(), self.get_attrs() + { + k: v + for k, v in self.get_variables().items() + if k in existing_variable_names + }, + self.get_attrs(), ) # Modified variables must use the same encoding as the store. vars_with_encoding = {} @@ -702,6 +707,17 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No if v.encoding == {"_FillValue": None} and fill_value is None: v.encoding = {} + # We need to do this for both new and existing variables to ensure we're not + # writing to a partial chunk, even though we don't use the `encoding` value + # when writing to an existing variable. See + # https://github.com/pydata/xarray/issues/8371 for details. + encoding = extract_zarr_variable_encoding( + v, + raise_on_invalid=check, + name=vn, + safe_chunks=self._safe_chunks, + ) + if name in existing_keys: # existing variable # TODO: if mode="a", consider overriding the existing variable @@ -732,9 +748,6 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No zarr_array = self.zarr_group[name] else: # new variable - encoding = extract_zarr_variable_encoding( - v, raise_on_invalid=check, name=vn, safe_chunks=self._safe_chunks - ) encoded_attrs = {} # the magic for storing the hidden dimension data encoded_attrs[DIMENSION_KEY] = dims diff --git a/xarray/coding/times.py b/xarray/coding/times.py index 92bce0abeaa..466e847e003 100644 --- a/xarray/coding/times.py +++ b/xarray/coding/times.py @@ -446,15 +446,7 @@ def format_cftime_datetime(date) -> str: """Converts a cftime.datetime object to a string with the format: YYYY-MM-DD HH:MM:SS.UUUUUU """ - return "{:04d}-{:02d}-{:02d} {:02d}:{:02d}:{:02d}.{:06d}".format( - date.year, - date.month, - date.day, - date.hour, - date.minute, - date.second, - date.microsecond, - ) + return f"{date.year:04d}-{date.month:02d}-{date.day:02d} {date.hour:02d}:{date.minute:02d}:{date.second:02d}.{date.microsecond:06d}" def infer_timedelta_units(deltas) -> str: diff --git a/xarray/coding/variables.py b/xarray/coding/variables.py index 52cf0fc3656..d31cb6e626a 100644 --- a/xarray/coding/variables.py +++ b/xarray/coding/variables.py @@ -81,9 +81,7 @@ def get_duck_array(self): return self.func(self.array.get_duck_array()) def __repr__(self) -> str: - return "{}({!r}, func={!r}, dtype={!r})".format( - type(self).__name__, self.array, self.func, self.dtype - ) + return f"{type(self).__name__}({self.array!r}, func={self.func!r}, dtype={self.dtype!r})" class NativeEndiannessArray(indexing.ExplicitlyIndexedNDArrayMixin): diff --git a/xarray/core/arithmetic.py b/xarray/core/arithmetic.py index 452c7115b75..734d7b328de 100644 --- a/xarray/core/arithmetic.py +++ b/xarray/core/arithmetic.py @@ -62,10 +62,10 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): if method != "__call__": # TODO: support other methods, e.g., reduce and accumulate. raise NotImplementedError( - "{} method for ufunc {} is not implemented on xarray objects, " + f"{method} method for ufunc {ufunc} is not implemented on xarray objects, " "which currently only support the __call__ method. As an " "alternative, consider explicitly converting xarray objects " - "to NumPy arrays (e.g., with `.values`).".format(method, ufunc) + "to NumPy arrays (e.g., with `.values`)." ) if any(isinstance(o, SupportsArithmetic) for o in out): diff --git a/xarray/core/computation.py b/xarray/core/computation.py index f29f6c4dd35..f09b04b7765 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -133,11 +133,7 @@ def __ne__(self, other): return not self == other def __repr__(self): - return "{}({!r}, {!r})".format( - type(self).__name__, - list(self.input_core_dims), - list(self.output_core_dims), - ) + return f"{type(self).__name__}({list(self.input_core_dims)!r}, {list(self.output_core_dims)!r})" def __str__(self): lhs = ",".join("({})".format(",".join(dims)) for dims in self.input_core_dims) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 389316d67c2..509962ff80d 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1126,6 +1126,8 @@ def load(self, **kwargs) -> Self: """Manually trigger loading of this array's data from disk or a remote source into memory and return this array. + Unlike compute, the original dataset is modified and returned. + Normally, it should not be necessary to call this method in user code, because all xarray functions should either work on deferred data or load data automatically. However, this method can be necessary when @@ -1148,8 +1150,9 @@ def load(self, **kwargs) -> Self: def compute(self, **kwargs) -> Self: """Manually trigger loading of this array's data from disk or a - remote source into memory and return a new array. The original is - left unaltered. + remote source into memory and return a new array. + + Unlike load, the original is left unaltered. Normally, it should not be necessary to call this method in user code, because all xarray functions should either work on deferred data or @@ -1161,6 +1164,11 @@ def compute(self, **kwargs) -> Self: **kwargs : dict Additional keyword arguments passed on to ``dask.compute``. + Returns + ------- + object : DataArray + New object with the data and all coordinates as in-memory arrays. + See Also -------- dask.compute @@ -1174,12 +1182,18 @@ def persist(self, **kwargs) -> Self: This keeps them as dask arrays but encourages them to keep data in memory. This is particularly useful when on a distributed machine. When on a single machine consider using ``.compute()`` instead. + Like compute (but unlike load), the original dataset is left unaltered. Parameters ---------- **kwargs : dict Additional keyword arguments passed on to ``dask.persist``. + Returns + ------- + object : DataArray + New object with all dask-backed data and coordinates as persisted dask arrays. + See Also -------- dask.persist @@ -4120,7 +4134,7 @@ def to_zarr( compute: Literal[True] = True, consolidated: bool | None = None, append_dim: Hashable | None = None, - region: Mapping[str, slice] | None = None, + region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None = None, safe_chunks: bool = True, storage_options: dict[str, str] | None = None, zarr_version: int | None = None, @@ -4140,7 +4154,7 @@ def to_zarr( compute: Literal[False], consolidated: bool | None = None, append_dim: Hashable | None = None, - region: Mapping[str, slice] | None = None, + region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None = None, safe_chunks: bool = True, storage_options: dict[str, str] | None = None, zarr_version: int | None = None, @@ -4158,7 +4172,7 @@ def to_zarr( compute: bool = True, consolidated: bool | None = None, append_dim: Hashable | None = None, - region: Mapping[str, slice] | None = None, + region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None = None, safe_chunks: bool = True, storage_options: dict[str, str] | None = None, zarr_version: int | None = None, @@ -4237,6 +4251,12 @@ def to_zarr( in with ``region``, use a separate call to ``to_zarr()`` with ``compute=False``. See "Appending to existing Zarr stores" in the reference documentation for full details. + + Users are expected to ensure that the specified region aligns with + Zarr chunk boundaries, and that dask chunks are also aligned. + Xarray makes limited checks that these multiple chunk boundaries line up. + It is possible to write incomplete chunks and corrupt the data with this + option if you are not careful. safe_chunks : bool, default: True If True, only allow writes to when there is a many-to-one relationship between Zarr chunks (specified in encoding) and Dask chunks. diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 4866bdbf988..96f3be00995 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -1006,6 +1006,11 @@ def compute(self, **kwargs) -> Self: **kwargs : dict Additional keyword arguments passed on to ``dask.compute``. + Returns + ------- + object : Dataset + New object with lazy data variables and coordinates as in-memory arrays. + See Also -------- dask.compute @@ -1038,12 +1043,18 @@ def persist(self, **kwargs) -> Self: operation keeps the data as dask arrays. This is particularly useful when using the dask.distributed scheduler and you want to load a large amount of data into distributed memory. + Like compute (but unlike load), the original dataset is left unaltered. Parameters ---------- **kwargs : dict Additional keyword arguments passed on to ``dask.persist``. + Returns + ------- + object : Dataset + New object with all dask-backed coordinates and data variables as persisted dask arrays. + See Also -------- dask.persist @@ -2453,6 +2464,12 @@ def to_zarr( in with ``region``, use a separate call to ``to_zarr()`` with ``compute=False``. See "Appending to existing Zarr stores" in the reference documentation for full details. + + Users are expected to ensure that the specified region aligns with + Zarr chunk boundaries, and that dask chunks are also aligned. + Xarray makes limited checks that these multiple chunk boundaries line up. + It is possible to write incomplete chunks and corrupt the data with this + option if you are not careful. safe_chunks : bool, default: True If True, only allow writes to when there is a many-to-one relationship between Zarr chunks (specified in encoding) and Dask chunks. diff --git a/xarray/core/formatting.py b/xarray/core/formatting.py index 260dabd9d31..3eed7d02a2e 100644 --- a/xarray/core/formatting.py +++ b/xarray/core/formatting.py @@ -289,8 +289,8 @@ def inline_sparse_repr(array): """Similar to sparse.COO.__repr__, but without the redundant shape/dtype.""" sparse_array_type = array_type("sparse") assert isinstance(array, sparse_array_type), array - return "<{}: nnz={:d}, fill_value={!s}>".format( - type(array).__name__, array.nnz, array.fill_value + return ( + f"<{type(array).__name__}: nnz={array.nnz:d}, fill_value={array.fill_value!s}>" ) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index e26c50c8b90..0926da6fd80 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -9,7 +9,7 @@ from dataclasses import dataclass, field from datetime import timedelta from html import escape -from typing import TYPE_CHECKING, Any, Callable +from typing import TYPE_CHECKING, Any, Callable, overload import numpy as np import pandas as pd @@ -236,14 +236,34 @@ def expanded_indexer(key, ndim): return tuple(new_key) -def _expand_slice(slice_, size: int) -> np.ndarray: - return np.arange(*slice_.indices(size)) +def _normalize_slice(sl: slice, size: int) -> slice: + """ + Ensure that given slice only contains positive start and stop values + (stop can be -1 for full-size slices with negative steps, e.g. [-10::-1]) + + Examples + -------- + >>> _normalize_slice(slice(0, 9), 10) + slice(0, 9, 1) + >>> _normalize_slice(slice(0, -1), 10) + slice(0, 9, 1) + """ + return slice(*sl.indices(size)) -def _normalize_slice(sl: slice, size) -> slice: - """Ensure that given slice only contains positive start and stop values - (stop can be -1 for full-size slices with negative steps, e.g. [-10::-1])""" - return slice(*sl.indices(size)) +def _expand_slice(slice_: slice, size: int) -> np.ndarray[Any, np.dtype[np.integer]]: + """ + Expand slice to an array containing only positive integers. + + Examples + -------- + >>> _expand_slice(slice(0, 9), 10) + array([0, 1, 2, 3, 4, 5, 6, 7, 8]) + >>> _expand_slice(slice(0, -1), 10) + array([0, 1, 2, 3, 4, 5, 6, 7, 8]) + """ + sl = _normalize_slice(slice_, size) + return np.arange(sl.start, sl.stop, sl.step) def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice: @@ -316,11 +336,15 @@ def __repr__(self) -> str: return f"{type(self).__name__}({self.tuple})" -def as_integer_or_none(value): +@overload +def as_integer_or_none(value: int) -> int: ... +@overload +def as_integer_or_none(value: None) -> None: ... +def as_integer_or_none(value: int | None) -> int | None: return None if value is None else operator.index(value) -def as_integer_slice(value): +def as_integer_slice(value: slice) -> slice: start = as_integer_or_none(value.start) stop = as_integer_or_none(value.stop) step = as_integer_or_none(value.step) diff --git a/xarray/core/merge.py b/xarray/core/merge.py index cbd06c8fdc5..a90e59e7c0b 100644 --- a/xarray/core/merge.py +++ b/xarray/core/merge.py @@ -562,25 +562,6 @@ def merge_coords( return variables, out_indexes -def assert_valid_explicit_coords( - variables: Mapping[Any, Any], - dims: Mapping[Any, int], - explicit_coords: Iterable[Hashable], -) -> None: - """Validate explicit coordinate names/dims. - - Raise a MergeError if an explicit coord shares a name with a dimension - but is comprised of arbitrary dimensions. - """ - for coord_name in explicit_coords: - if coord_name in dims and variables[coord_name].dims != (coord_name,): - raise MergeError( - f"coordinate {coord_name} shares a name with a dataset dimension, but is " - "not a 1D variable along that dimension. This is disallowed " - "by the xarray data model." - ) - - def merge_attrs(variable_attrs, combine_attrs, context=None): """Combine attributes from different variables according to combine_attrs""" if not variable_attrs: @@ -728,7 +709,6 @@ def merge_core( # coordinates may be dropped in merged results coord_names.intersection_update(variables) if explicit_coords is not None: - assert_valid_explicit_coords(variables, dims, explicit_coords) coord_names.update(explicit_coords) for dim, size in dims.items(): if dim in variables: diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 14243487d85..2229eaa2d24 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -126,13 +126,18 @@ def as_variable( if isinstance(obj, Variable): obj = obj.copy(deep=False) elif isinstance(obj, tuple): - if isinstance(obj[1], DataArray): + try: + dims_, data_, *attrs = obj + except ValueError: + raise ValueError(f"Tuple {obj} is not in the form (dims, data[, attrs])") + + if isinstance(data_, DataArray): raise TypeError( f"Variable {name!r}: Using a DataArray object to construct a variable is" " ambiguous, please extract the data using the .data property." ) try: - obj = Variable(*obj) + obj = Variable(dims_, data_, *attrs) except (TypeError, ValueError) as error: raise error.__class__( f"Variable {name!r}: Could not convert tuple of form " @@ -269,8 +274,13 @@ def as_compatible_data( from xarray.core.dataarray import DataArray - if isinstance(data, (Variable, DataArray)): - return data.data + # TODO: do this uwrapping in the Variable/NamedArray constructor instead. + if isinstance(data, Variable): + return cast("T_DuckArray", data._data) + + # TODO: do this uwrapping in the DataArray constructor instead. + if isinstance(data, DataArray): + return cast("T_DuckArray", data._variable._data) if isinstance(data, NON_NUMPY_SUPPORTED_ARRAY_TYPES): data = _possibly_convert_datetime_or_timedelta_index(data) diff --git a/xarray/core/weighted.py b/xarray/core/weighted.py index ae9521309e0..8cb90ac1b2b 100644 --- a/xarray/core/weighted.py +++ b/xarray/core/weighted.py @@ -84,7 +84,7 @@ method supported by this weighted version corresponds to the default "linear" option of ``numpy.quantile``. This is "Type 7" option, described in Hyndman and Fan (1996) [2]_. The implementation is largely inspired by a blog post - from A. Akinshin's [3]_. + from A. Akinshin's (2023) [3]_. Parameters ---------- @@ -122,7 +122,8 @@ .. [1] https://notstatschat.rbind.io/2020/08/04/weights-in-statistics/ .. [2] Hyndman, R. J. & Fan, Y. (1996). Sample Quantiles in Statistical Packages. The American Statistician, 50(4), 361–365. https://doi.org/10.2307/2684934 - .. [3] https://aakinshin.net/posts/weighted-quantiles + .. [3] Akinshin, A. (2023) "Weighted quantile estimators" arXiv:2304.07265 [stat.ME] + https://arxiv.org/abs/2304.07265 """ diff --git a/xarray/plot/dataarray_plot.py b/xarray/plot/dataarray_plot.py index 8386161bf29..ed752d3461f 100644 --- a/xarray/plot/dataarray_plot.py +++ b/xarray/plot/dataarray_plot.py @@ -1848,9 +1848,10 @@ def _center_pixels(x): # missing data transparent. We therefore add an alpha channel if # there isn't one, and set it to transparent where data is masked. if z.shape[-1] == 3: - alpha = np.ma.ones(z.shape[:2] + (1,), dtype=z.dtype) + safe_dtype = np.promote_types(z.dtype, np.uint8) + alpha = np.ma.ones(z.shape[:2] + (1,), dtype=safe_dtype) if np.issubdtype(z.dtype, np.integer): - alpha *= 255 + alpha[:] = 255 z = np.ma.concatenate((z, alpha), axis=2) else: z = z.copy() diff --git a/xarray/plot/utils.py b/xarray/plot/utils.py index 804e1cfd795..8789bc2f9c2 100644 --- a/xarray/plot/utils.py +++ b/xarray/plot/utils.py @@ -4,7 +4,7 @@ import textwrap import warnings from collections.abc import Hashable, Iterable, Mapping, MutableMapping, Sequence -from datetime import datetime +from datetime import date, datetime from inspect import getfullargspec from typing import TYPE_CHECKING, Any, Callable, Literal, overload @@ -672,7 +672,7 @@ def _ensure_plottable(*args) -> None: np.bool_, np.str_, ) - other_types: tuple[type[object], ...] = (datetime,) + other_types: tuple[type[object], ...] = (datetime, date) cftime_datetime_types: tuple[type[object], ...] = ( () if cftime is None else (cftime.datetime,) ) diff --git a/xarray/testing/strategies.py b/xarray/testing/strategies.py index 79385bc5226..449d0c793cc 100644 --- a/xarray/testing/strategies.py +++ b/xarray/testing/strategies.py @@ -21,6 +21,7 @@ __all__ = [ "supported_dtypes", + "pandas_index_dtypes", "names", "dimension_names", "dimension_sizes", @@ -59,6 +60,26 @@ def supported_dtypes() -> st.SearchStrategy[np.dtype]: | npst.unsigned_integer_dtypes(endianness="=") | npst.floating_dtypes(endianness="=") | npst.complex_number_dtypes(endianness="=") + # | npst.datetime64_dtypes() + # | npst.timedelta64_dtypes() + # | npst.unicode_string_dtypes() + ) + + +def pandas_index_dtypes() -> st.SearchStrategy[np.dtype]: + """ + Dtypes supported by pandas indexes. + Restrict datetime64 and timedelta64 to ns frequency till Xarray relaxes that. + """ + return ( + npst.integer_dtypes(endianness="=", sizes=(32, 64)) + | npst.unsigned_integer_dtypes(endianness="=", sizes=(32, 64)) + | npst.floating_dtypes(endianness="=", sizes=(32, 64)) + # TODO: unset max_period + | npst.datetime64_dtypes(endianness="=", max_period="ns") + # TODO: set max_period="D" + | npst.timedelta64_dtypes(endianness="=", max_period="ns") + | npst.unicode_string_dtypes(endianness="=") ) @@ -87,6 +108,7 @@ def names() -> st.SearchStrategy[str]: def dimension_names( *, + name_strategy=names(), min_dims: int = 0, max_dims: int = 3, ) -> st.SearchStrategy[list[Hashable]]: @@ -97,6 +119,8 @@ def dimension_names( Parameters ---------- + name_strategy + Strategy for making names. Useful if we need to share this. min_dims Minimum number of dimensions in generated list. max_dims @@ -104,7 +128,7 @@ def dimension_names( """ return st.lists( - elements=names(), + elements=name_strategy, min_size=min_dims, max_size=max_dims, unique=True, diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 1685e2f4824..3ce788dfb7f 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -115,7 +115,6 @@ def _importorskip( has_fsspec, requires_fsspec = _importorskip("fsspec") has_iris, requires_iris = _importorskip("iris") has_numbagg, requires_numbagg = _importorskip("numbagg", "0.4.0") -has_plum, requires_plum = _importorskip("plum") has_pyarrow, requires_pyarrow = _importorskip("pyarrow") with warnings.catch_warnings(): warnings.filterwarnings( diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 07573066568..be9b3ef0422 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -2261,7 +2261,6 @@ def test_write_uneven_dask_chunks(self) -> None: original = create_test_data().chunk({"dim1": 3, "dim2": 4, "dim3": 3}) with self.roundtrip(original, open_kwargs={"chunks": {}}) as actual: for k, v in actual.data_vars.items(): - print(k) assert v.chunks == actual[k].chunks def test_chunk_encoding(self) -> None: @@ -2305,7 +2304,7 @@ def test_chunk_encoding_with_dask(self) -> None: # should fail if encoding["chunks"] clashes with dask_chunks badenc = ds.chunk({"x": 4}) badenc.var1.encoding["chunks"] = (6,) - with pytest.raises(NotImplementedError, match=r"named 'var1' would overlap"): + with pytest.raises(ValueError, match=r"named 'var1' would overlap"): with self.roundtrip(badenc) as actual: pass @@ -2343,9 +2342,7 @@ def test_chunk_encoding_with_dask(self) -> None: # but itermediate unaligned chunks are bad badenc = ds.chunk({"x": (3, 5, 3, 1)}) badenc.var1.encoding["chunks"] = (3,) - with pytest.raises( - NotImplementedError, match=r"would overlap multiple dask chunks" - ): + with pytest.raises(ValueError, match=r"would overlap multiple dask chunks"): with self.roundtrip(badenc) as actual: pass @@ -2359,7 +2356,7 @@ def test_chunk_encoding_with_dask(self) -> None: # TODO: remove this failure once synchronized overlapping writes are # supported by xarray ds_chunk4["var1"].encoding.update({"chunks": 5}) - with pytest.raises(NotImplementedError, match=r"named 'var1' would overlap"): + with pytest.raises(ValueError, match=r"named 'var1' would overlap"): with self.roundtrip(ds_chunk4) as actual: pass # override option @@ -2468,6 +2465,27 @@ def test_group(self) -> None: ) as actual: assert_identical(original, actual) + def test_zarr_mode_w_overwrites_encoding(self) -> None: + import zarr + + data = Dataset({"foo": ("x", [1.0, 1.0, 1.0])}) + with self.create_zarr_target() as store: + data.to_zarr( + store, **self.version_kwargs, encoding={"foo": {"add_offset": 1}} + ) + np.testing.assert_equal( + zarr.open_group(store, **self.version_kwargs)["foo"], data.foo.data - 1 + ) + data.to_zarr( + store, + **self.version_kwargs, + encoding={"foo": {"add_offset": 0}}, + mode="w", + ) + np.testing.assert_equal( + zarr.open_group(store, **self.version_kwargs)["foo"], data.foo.data + ) + def test_encoding_kwarg_fixed_width_string(self) -> None: # not relevant for zarr, since we don't use EncodedStringCoder pass @@ -5733,3 +5751,80 @@ def test_zarr_region(tmp_path): # Write without region ds_transposed.to_zarr(tmp_path / "test.zarr", mode="r+") + + +@requires_zarr +@requires_dask +def test_zarr_region_chunk_partial(tmp_path): + """ + Check that writing to partial chunks with `region` fails, assuming `safe_chunks=False`. + """ + ds = ( + xr.DataArray(np.arange(120).reshape(4, 3, -1), dims=list("abc")) + .rename("var1") + .to_dataset() + ) + + ds.chunk(5).to_zarr(tmp_path / "foo.zarr", compute=False, mode="w") + with pytest.raises(ValueError): + for r in range(ds.sizes["a"]): + ds.chunk(3).isel(a=[r]).to_zarr( + tmp_path / "foo.zarr", region=dict(a=slice(r, r + 1)) + ) + + +@requires_zarr +@requires_dask +def test_zarr_append_chunk_partial(tmp_path): + t_coords = np.array([np.datetime64("2020-01-01").astype("datetime64[ns]")]) + data = np.ones((10, 10)) + + da = xr.DataArray( + data.reshape((-1, 10, 10)), + dims=["time", "x", "y"], + coords={"time": t_coords}, + name="foo", + ) + da.to_zarr(tmp_path / "foo.zarr", mode="w", encoding={"foo": {"chunks": (5, 5, 1)}}) + + new_time = np.array([np.datetime64("2021-01-01").astype("datetime64[ns]")]) + + da2 = xr.DataArray( + data.reshape((-1, 10, 10)), + dims=["time", "x", "y"], + coords={"time": new_time}, + name="foo", + ) + with pytest.raises(ValueError, match="encoding was provided"): + da2.to_zarr( + tmp_path / "foo.zarr", + append_dim="time", + mode="a", + encoding={"foo": {"chunks": (1, 1, 1)}}, + ) + + # chunking with dask sidesteps the encoding check, so we need a different check + with pytest.raises(ValueError, match="Specified zarr chunks"): + da2.chunk({"x": 1, "y": 1, "time": 1}).to_zarr( + tmp_path / "foo.zarr", append_dim="time", mode="a" + ) + + +@requires_zarr +@requires_dask +def test_zarr_region_chunk_partial_offset(tmp_path): + # https://github.com/pydata/xarray/pull/8459#issuecomment-1819417545 + store = tmp_path / "foo.zarr" + data = np.ones((30,)) + da = xr.DataArray(data, dims=["x"], coords={"x": range(30)}, name="foo").chunk(x=10) + da.to_zarr(store, compute=False) + + da.isel(x=slice(10)).chunk(x=(10,)).to_zarr(store, region="auto") + + da.isel(x=slice(5, 25)).chunk(x=(10, 10)).to_zarr( + store, safe_chunks=False, region="auto" + ) + + # This write is unsafe, and should raise an error, but does not. + # with pytest.raises(ValueError): + # da.isel(x=slice(5, 25)).chunk(x=(10, 10)).to_zarr(store, region="auto") diff --git a/xarray/tests/test_coordinates.py b/xarray/tests/test_coordinates.py index 40743194ce6..f88e554d333 100644 --- a/xarray/tests/test_coordinates.py +++ b/xarray/tests/test_coordinates.py @@ -1,5 +1,6 @@ from __future__ import annotations +import numpy as np import pandas as pd import pytest @@ -8,7 +9,7 @@ from xarray.core.dataarray import DataArray from xarray.core.dataset import Dataset from xarray.core.indexes import PandasIndex, PandasMultiIndex -from xarray.core.variable import IndexVariable +from xarray.core.variable import IndexVariable, Variable from xarray.tests import assert_identical, source_ndarray @@ -174,3 +175,10 @@ def test_align(self) -> None: left2, right2 = align(left, right, join="override") assert_identical(left2, left) assert_identical(left2, right2) + + def test_dataset_from_coords_with_multidim_var_same_name(self): + # regression test for GH #8883 + var = Variable(data=np.arange(6).reshape(2, 3), dims=["x", "y"]) + coords = Coordinates(coords={"x": var}, indexes={}) + ds = Dataset(coords=coords) + assert ds.coords["x"].dims == ("x", "y") diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 04112a16ab3..f3413ea40a4 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -4947,7 +4947,7 @@ def test_idxmin( with pytest.raises(ValueError): xr.DataArray(5).idxmin() - coordarr0 = xr.DataArray(ar0.coords["x"], dims=["x"]) + coordarr0 = xr.DataArray(ar0.coords["x"].data, dims=["x"]) coordarr1 = coordarr0.copy() hasna = np.isnan(minindex) @@ -5062,7 +5062,7 @@ def test_idxmax( with pytest.raises(ValueError): xr.DataArray(5).idxmax() - coordarr0 = xr.DataArray(ar0.coords["x"], dims=["x"]) + coordarr0 = xr.DataArray(ar0.coords["x"].data, dims=["x"]) coordarr1 = coordarr0.copy() hasna = np.isnan(maxindex) @@ -7167,3 +7167,13 @@ def test_nD_coord_dataarray() -> None: _assert_internal_invariants(da4, check_default_indexes=True) assert "x" not in da4.xindexes assert "x" in da4.coords + + +def test_lazy_data_variable_not_loaded(): + # GH8753 + array = InaccessibleArray(np.array([1, 2, 3])) + v = Variable(data=array, dims="x") + # No data needs to be accessed, so no error should be raised + da = xr.DataArray(v) + # No data needs to be accessed, so no error should be raised + xr.DataArray(da) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index f85f3f36d38..c5f9d6b157c 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -486,14 +486,6 @@ def test_constructor(self) -> None: actual = Dataset({"z": expected["z"]}) assert_identical(expected, actual) - def test_constructor_invalid_dims(self) -> None: - # regression for GH1120 - with pytest.raises(MergeError): - Dataset( - data_vars=dict(v=("y", [1, 2, 3, 4])), - coords=dict(y=DataArray([0.1, 0.2, 0.3, 0.4], dims="x")), - ) - def test_constructor_1d(self) -> None: expected = Dataset({"x": (["x"], 5.0 + np.arange(5))}) actual = Dataset({"x": 5.0 + np.arange(5)}) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index ee80363970b..afe4d669628 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -13,10 +13,10 @@ import xarray as xr from xarray import DataArray, Dataset, Variable from xarray.core.groupby import _consolidate_slices +from xarray.core.types import InterpOptions from xarray.tests import ( InaccessibleArray, assert_allclose, - assert_array_equal, assert_equal, assert_identical, create_test_data, @@ -30,7 +30,7 @@ @pytest.fixture -def dataset(): +def dataset() -> xr.Dataset: ds = xr.Dataset( { "foo": (("x", "y", "z"), np.random.randn(3, 4, 2)), @@ -45,7 +45,7 @@ def dataset(): @pytest.fixture -def array(dataset): +def array(dataset) -> xr.DataArray: return dataset["foo"] @@ -525,7 +525,7 @@ def test_ds_groupby_quantile() -> None: @pytest.mark.parametrize("as_dataset", [False, True]) -def test_groupby_quantile_interpolation_deprecated(as_dataset) -> None: +def test_groupby_quantile_interpolation_deprecated(as_dataset: bool) -> None: array = xr.DataArray(data=[1, 2, 3, 4], coords={"x": [1, 1, 2, 2]}, dims="x") arr: xr.DataArray | xr.Dataset @@ -904,7 +904,7 @@ def test_groupby_dataset_reduce() -> None: @pytest.mark.parametrize("squeeze", [True, False]) -def test_groupby_dataset_math(squeeze) -> None: +def test_groupby_dataset_math(squeeze: bool) -> None: def reorder_dims(x): return x.transpose("dim1", "dim2", "dim3", "time") @@ -1125,7 +1125,7 @@ def test_groupby_dataset_order() -> None: # .assertEqual(all_vars, all_vars_ref) -def test_groupby_dataset_fillna(): +def test_groupby_dataset_fillna() -> None: ds = Dataset({"a": ("x", [np.nan, 1, np.nan, 3])}, {"x": [0, 1, 2, 3]}) expected = Dataset({"a": ("x", range(4))}, {"x": [0, 1, 2, 3]}) for target in [ds, expected]: @@ -1145,12 +1145,12 @@ def test_groupby_dataset_fillna(): assert actual.a.attrs == ds.a.attrs -def test_groupby_dataset_where(): +def test_groupby_dataset_where() -> None: # groupby ds = Dataset({"a": ("x", range(5))}, {"c": ("x", [0, 0, 1, 1, 1])}) cond = Dataset({"a": ("c", [True, False])}) expected = ds.copy(deep=True) - expected["a"].values = [0, 1] + [np.nan] * 3 + expected["a"].values = np.array([0, 1] + [np.nan] * 3) actual = ds.groupby("c").where(cond) assert_identical(expected, actual) @@ -1163,7 +1163,7 @@ def test_groupby_dataset_where(): assert actual.a.attrs == ds.a.attrs -def test_groupby_dataset_assign(): +def test_groupby_dataset_assign() -> None: ds = Dataset({"a": ("x", range(3))}, {"b": ("x", ["A"] * 2 + ["B"])}) actual = ds.groupby("b").assign(c=lambda ds: 2 * ds.a) expected = ds.merge({"c": ("x", [0, 2, 4])}) @@ -1178,7 +1178,7 @@ def test_groupby_dataset_assign(): assert_identical(actual, expected) -def test_groupby_dataset_map_dataarray_func(): +def test_groupby_dataset_map_dataarray_func() -> None: # regression GH6379 ds = Dataset({"foo": ("x", [1, 2, 3, 4])}, coords={"x": [0, 0, 1, 1]}) actual = ds.groupby("x").map(lambda grp: grp.foo.mean()) @@ -1186,7 +1186,7 @@ def test_groupby_dataset_map_dataarray_func(): assert_identical(actual, expected) -def test_groupby_dataarray_map_dataset_func(): +def test_groupby_dataarray_map_dataset_func() -> None: # regression GH6379 da = DataArray([1, 2, 3, 4], coords={"x": [0, 0, 1, 1]}, dims="x", name="foo") actual = da.groupby("x").map(lambda grp: grp.mean().to_dataset()) @@ -1196,7 +1196,7 @@ def test_groupby_dataarray_map_dataset_func(): @requires_flox @pytest.mark.parametrize("kwargs", [{"method": "map-reduce"}, {"engine": "numpy"}]) -def test_groupby_flox_kwargs(kwargs): +def test_groupby_flox_kwargs(kwargs) -> None: ds = Dataset({"a": ("x", range(5))}, {"c": ("x", [0, 0, 1, 1, 1])}) with xr.set_options(use_flox=False): expected = ds.groupby("c").mean() @@ -1207,7 +1207,7 @@ def test_groupby_flox_kwargs(kwargs): class TestDataArrayGroupBy: @pytest.fixture(autouse=True) - def setup(self): + def setup(self) -> None: self.attrs = {"attr1": "value1", "attr2": 2929} self.x = np.random.random((10, 20)) self.v = Variable(["x", "y"], self.x) @@ -1224,7 +1224,7 @@ def setup(self): self.da.coords["abc"] = ("y", np.array(["a"] * 9 + ["c"] + ["b"] * 10)) self.da.coords["y"] = 20 + 100 * self.da["y"] - def test_stack_groupby_unsorted_coord(self): + def test_stack_groupby_unsorted_coord(self) -> None: data = [[0, 1], [2, 3]] data_flat = [0, 1, 2, 3] dims = ["x", "y"] @@ -1243,7 +1243,7 @@ def test_stack_groupby_unsorted_coord(self): expected2 = xr.DataArray(data_flat, dims=["z"], coords={"z": midx2}) assert_equal(actual2, expected2) - def test_groupby_iter(self): + def test_groupby_iter(self) -> None: for (act_x, act_dv), (exp_x, exp_ds) in zip( self.dv.groupby("y", squeeze=False), self.ds.groupby("y", squeeze=False) ): @@ -1255,12 +1255,19 @@ def test_groupby_iter(self): ): assert_identical(exp_dv, act_dv) - def test_groupby_properties(self): + def test_groupby_properties(self) -> None: grouped = self.da.groupby("abc") expected_groups = {"a": range(0, 9), "c": [9], "b": range(10, 20)} assert expected_groups.keys() == grouped.groups.keys() for key in expected_groups: - assert_array_equal(expected_groups[key], grouped.groups[key]) + expected_group = expected_groups[key] + actual_group = grouped.groups[key] + + # TODO: array_api doesn't allow slice: + assert not isinstance(expected_group, slice) + assert not isinstance(actual_group, slice) + + np.testing.assert_array_equal(expected_group, actual_group) assert 3 == len(grouped) @pytest.mark.parametrize( @@ -1284,7 +1291,7 @@ def identity(x): if (by.name if use_da else by) != "abc": assert len(recwarn) == (1 if squeeze in [None, True] else 0) - def test_groupby_sum(self): + def test_groupby_sum(self) -> None: array = self.da grouped = array.groupby("abc") @@ -1338,7 +1345,7 @@ def test_groupby_sum(self): assert_allclose(expected_sum_axis1, grouped.sum("y")) @pytest.mark.parametrize("method", ["sum", "mean", "median"]) - def test_groupby_reductions(self, method): + def test_groupby_reductions(self, method) -> None: array = self.da grouped = array.groupby("abc") @@ -1368,7 +1375,7 @@ def test_groupby_reductions(self, method): assert_allclose(expected, actual_legacy) assert_allclose(expected, actual_npg) - def test_groupby_count(self): + def test_groupby_count(self) -> None: array = DataArray( [0, 0, np.nan, np.nan, 0, 0], coords={"cat": ("x", ["a", "b", "b", "c", "c", "c"])}, @@ -1380,7 +1387,9 @@ def test_groupby_count(self): @pytest.mark.parametrize("shortcut", [True, False]) @pytest.mark.parametrize("keep_attrs", [None, True, False]) - def test_groupby_reduce_keep_attrs(self, shortcut, keep_attrs): + def test_groupby_reduce_keep_attrs( + self, shortcut: bool, keep_attrs: bool | None + ) -> None: array = self.da array.attrs["foo"] = "bar" @@ -1392,7 +1401,7 @@ def test_groupby_reduce_keep_attrs(self, shortcut, keep_attrs): assert_identical(expected, actual) @pytest.mark.parametrize("keep_attrs", [None, True, False]) - def test_groupby_keep_attrs(self, keep_attrs): + def test_groupby_keep_attrs(self, keep_attrs: bool | None) -> None: array = self.da array.attrs["foo"] = "bar" @@ -1406,7 +1415,7 @@ def test_groupby_keep_attrs(self, keep_attrs): actual.data = expected.data assert_identical(expected, actual) - def test_groupby_map_center(self): + def test_groupby_map_center(self) -> None: def center(x): return x - np.mean(x) @@ -1421,14 +1430,14 @@ def center(x): expected_centered = expected_ds["foo"] assert_allclose(expected_centered, grouped.map(center)) - def test_groupby_map_ndarray(self): + def test_groupby_map_ndarray(self) -> None: # regression test for #326 array = self.da grouped = array.groupby("abc") - actual = grouped.map(np.asarray) + actual = grouped.map(np.asarray) # type: ignore[arg-type] # TODO: Not sure using np.asarray like this makes sense with array api assert_equal(array, actual) - def test_groupby_map_changes_metadata(self): + def test_groupby_map_changes_metadata(self) -> None: def change_metadata(x): x.coords["x"] = x.coords["x"] * 2 x.attrs["fruit"] = "lemon" @@ -1442,7 +1451,7 @@ def change_metadata(x): assert_equal(expected, actual) @pytest.mark.parametrize("squeeze", [True, False]) - def test_groupby_math_squeeze(self, squeeze): + def test_groupby_math_squeeze(self, squeeze: bool) -> None: array = self.da grouped = array.groupby("x", squeeze=squeeze) @@ -1461,7 +1470,7 @@ def test_groupby_math_squeeze(self, squeeze): actual = ds + grouped assert_identical(expected, actual) - def test_groupby_math(self): + def test_groupby_math(self) -> None: array = self.da grouped = array.groupby("abc") expected_agg = (grouped.mean(...) - np.arange(3)).rename(None) @@ -1470,13 +1479,13 @@ def test_groupby_math(self): assert_allclose(expected_agg, actual_agg) with pytest.raises(TypeError, match=r"only support binary ops"): - grouped + 1 + grouped + 1 # type: ignore[type-var] with pytest.raises(TypeError, match=r"only support binary ops"): - grouped + grouped + grouped + grouped # type: ignore[type-var] with pytest.raises(TypeError, match=r"in-place operations"): - array += grouped + array += grouped # type: ignore[arg-type] - def test_groupby_math_not_aligned(self): + def test_groupby_math_not_aligned(self) -> None: array = DataArray( range(4), {"b": ("x", [0, 0, 1, 1]), "x": [0, 1, 2, 3]}, dims="x" ) @@ -1497,12 +1506,12 @@ def test_groupby_math_not_aligned(self): expected.coords["c"] = (["x"], [123] * 2 + [np.nan] * 2) assert_identical(expected, actual) - other = Dataset({"a": ("b", [10])}, {"b": [0]}) - actual = array.groupby("b") + other - expected = Dataset({"a": ("x", [10, 11, np.nan, np.nan])}, array.coords) - assert_identical(expected, actual) + other_ds = Dataset({"a": ("b", [10])}, {"b": [0]}) + actual_ds = array.groupby("b") + other_ds + expected_ds = Dataset({"a": ("x", [10, 11, np.nan, np.nan])}, array.coords) + assert_identical(expected_ds, actual_ds) - def test_groupby_restore_dim_order(self): + def test_groupby_restore_dim_order(self) -> None: array = DataArray( np.random.randn(5, 3), coords={"a": ("x", range(5)), "b": ("y", range(3))}, @@ -1517,7 +1526,7 @@ def test_groupby_restore_dim_order(self): result = array.groupby(by, squeeze=False).map(lambda x: x.squeeze()) assert result.dims == expected_dims - def test_groupby_restore_coord_dims(self): + def test_groupby_restore_coord_dims(self) -> None: array = DataArray( np.random.randn(5, 3), coords={ @@ -1539,7 +1548,7 @@ def test_groupby_restore_coord_dims(self): )["c"] assert result.dims == expected_dims - def test_groupby_first_and_last(self): + def test_groupby_first_and_last(self) -> None: array = DataArray([1, 2, 3, 4, 5], dims="x") by = DataArray(["a"] * 2 + ["b"] * 3, dims="x", name="ab") @@ -1560,7 +1569,7 @@ def test_groupby_first_and_last(self): expected = array # should be a no-op assert_identical(expected, actual) - def make_groupby_multidim_example_array(self): + def make_groupby_multidim_example_array(self) -> DataArray: return DataArray( [[[0, 1], [2, 3]], [[5, 10], [15, 20]]], coords={ @@ -1570,7 +1579,7 @@ def make_groupby_multidim_example_array(self): dims=["time", "ny", "nx"], ) - def test_groupby_multidim(self): + def test_groupby_multidim(self) -> None: array = self.make_groupby_multidim_example_array() for dim, expected_sum in [ ("lon", DataArray([5, 28, 23], coords=[("lon", [30.0, 40.0, 50.0])])), @@ -1579,7 +1588,7 @@ def test_groupby_multidim(self): actual_sum = array.groupby(dim).sum(...) assert_identical(expected_sum, actual_sum) - def test_groupby_multidim_map(self): + def test_groupby_multidim_map(self) -> None: array = self.make_groupby_multidim_example_array() actual = array.groupby("lon").map(lambda x: x - x.mean()) expected = DataArray( @@ -1640,7 +1649,7 @@ def test_groupby_bins( # make sure original array dims are unchanged assert len(array.dim_0) == 4 - def test_groupby_bins_ellipsis(self): + def test_groupby_bins_ellipsis(self) -> None: da = xr.DataArray(np.ones((2, 3, 4))) bins = [-1, 0, 1, 2] with xr.set_options(use_flox=False): @@ -1679,7 +1688,7 @@ def test_groupby_bins_gives_correct_subset(self, use_flox: bool) -> None: actual = gb.count() assert_identical(actual, expected) - def test_groupby_bins_empty(self): + def test_groupby_bins_empty(self) -> None: array = DataArray(np.arange(4), [("x", range(4))]) # one of these bins will be empty bins = [0, 4, 5] @@ -1691,7 +1700,7 @@ def test_groupby_bins_empty(self): # (was a problem in earlier versions) assert len(array.x) == 4 - def test_groupby_bins_multidim(self): + def test_groupby_bins_multidim(self) -> None: array = self.make_groupby_multidim_example_array() bins = [0, 15, 20] bin_coords = pd.cut(array["lat"].values.flat, bins).categories @@ -1725,7 +1734,7 @@ def test_groupby_bins_multidim(self): ) assert_identical(actual, expected) - def test_groupby_bins_sort(self): + def test_groupby_bins_sort(self) -> None: data = xr.DataArray( np.arange(100), dims="x", coords={"x": np.linspace(-100, 100, num=100)} ) @@ -1738,14 +1747,14 @@ def test_groupby_bins_sort(self): expected = data.groupby_bins("x", bins=11).count() assert_identical(actual, expected) - def test_groupby_assign_coords(self): + def test_groupby_assign_coords(self) -> None: array = DataArray([1, 2, 3, 4], {"c": ("x", [0, 0, 1, 1])}, dims="x") actual = array.groupby("c").assign_coords(d=lambda a: a.mean()) expected = array.copy() expected.coords["d"] = ("x", [1.5, 1.5, 3.5, 3.5]) assert_identical(actual, expected) - def test_groupby_fillna(self): + def test_groupby_fillna(self) -> None: a = DataArray([np.nan, 1, np.nan, 3], coords={"x": range(4)}, dims="x") fill_value = DataArray([0, 1], dims="y") actual = a.fillna(fill_value) @@ -1831,7 +1840,7 @@ def test_resample_doctest(self, use_cftime: bool) -> None: ) assert_identical(actual, expected) - def test_da_resample_func_args(self): + def test_da_resample_func_args(self) -> None: def func(arg1, arg2, arg3=0.0): return arg1.mean("time") + arg2 + arg3 @@ -1841,7 +1850,7 @@ def func(arg1, arg2, arg3=0.0): actual = da.resample(time="D").map(func, args=(1.0,), arg3=1.0) assert_identical(actual, expected) - def test_resample_first(self): + def test_resample_first(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.arange(10), [("time", times)]) @@ -1878,14 +1887,14 @@ def test_resample_first(self): expected = DataArray(expected_times, [("time", times[::4])], name="time") assert_identical(expected, actual) - def test_resample_bad_resample_dim(self): + def test_resample_bad_resample_dim(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.arange(10), [("__resample_dim__", times)]) with pytest.raises(ValueError, match=r"Proxy resampling dimension"): - array.resample(**{"__resample_dim__": "1D"}).first() + array.resample(**{"__resample_dim__": "1D"}).first() # type: ignore[arg-type] @requires_scipy - def test_resample_drop_nondim_coords(self): + def test_resample_drop_nondim_coords(self) -> None: xs = np.arange(6) ys = np.arange(3) times = pd.date_range("2000-01-01", freq="6h", periods=5) @@ -1916,7 +1925,7 @@ def test_resample_drop_nondim_coords(self): ) assert "tc" not in actual.coords - def test_resample_keep_attrs(self): + def test_resample_keep_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.ones(10), [("time", times)]) array.attrs["meta"] = "data" @@ -1925,7 +1934,7 @@ def test_resample_keep_attrs(self): expected = DataArray([1, 1, 1], [("time", times[::4])], attrs=array.attrs) assert_identical(result, expected) - def test_resample_skipna(self): + def test_resample_skipna(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) array = DataArray(np.ones(10), [("time", times)]) array[1] = np.nan @@ -1934,7 +1943,7 @@ def test_resample_skipna(self): expected = DataArray([np.nan, 1, 1], [("time", times[::4])]) assert_identical(result, expected) - def test_upsample(self): + def test_upsample(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=5) array = DataArray(np.arange(5), [("time", times)]) @@ -1965,7 +1974,7 @@ def test_upsample(self): expected = DataArray(array.reindex(time=new_times, method="nearest")) assert_identical(expected, actual) - def test_upsample_nd(self): + def test_upsample_nd(self) -> None: # Same as before, but now we try on multi-dimensional DataArrays. xs = np.arange(6) ys = np.arange(3) @@ -2023,29 +2032,29 @@ def test_upsample_nd(self): ) assert_identical(expected, actual) - def test_upsample_tolerance(self): + def test_upsample_tolerance(self) -> None: # Test tolerance keyword for upsample methods bfill, pad, nearest times = pd.date_range("2000-01-01", freq="1D", periods=2) times_upsampled = pd.date_range("2000-01-01", freq="6h", periods=5) array = DataArray(np.arange(2), [("time", times)]) # Forward fill - actual = array.resample(time="6h").ffill(tolerance="12h") + actual = array.resample(time="6h").ffill(tolerance="12h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. expected = DataArray([0.0, 0.0, 0.0, np.nan, 1.0], [("time", times_upsampled)]) assert_identical(expected, actual) # Backward fill - actual = array.resample(time="6h").bfill(tolerance="12h") + actual = array.resample(time="6h").bfill(tolerance="12h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. expected = DataArray([0.0, np.nan, 1.0, 1.0, 1.0], [("time", times_upsampled)]) assert_identical(expected, actual) # Nearest - actual = array.resample(time="6h").nearest(tolerance="6h") + actual = array.resample(time="6h").nearest(tolerance="6h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. expected = DataArray([0, 0, np.nan, 1, 1], [("time", times_upsampled)]) assert_identical(expected, actual) @requires_scipy - def test_upsample_interpolate(self): + def test_upsample_interpolate(self) -> None: from scipy.interpolate import interp1d xs = np.arange(6) @@ -2060,7 +2069,15 @@ def test_upsample_interpolate(self): # Split the times into equal sub-intervals to simulate the 6 hour # to 1 hour up-sampling new_times_idx = np.linspace(0, len(times) - 1, len(times) * 5) - for kind in ["linear", "nearest", "zero", "slinear", "quadratic", "cubic"]: + kinds: list[InterpOptions] = [ + "linear", + "nearest", + "zero", + "slinear", + "quadratic", + "cubic", + ] + for kind in kinds: actual = array.resample(time="1h").interpolate(kind) f = interp1d( np.arange(len(times)), @@ -2083,7 +2100,7 @@ def test_upsample_interpolate(self): @requires_scipy @pytest.mark.filterwarnings("ignore:Converting non-nanosecond") - def test_upsample_interpolate_bug_2197(self): + def test_upsample_interpolate_bug_2197(self) -> None: dates = pd.date_range("2007-02-01", "2007-03-01", freq="D") da = xr.DataArray(np.arange(len(dates)), [("time", dates)]) result = da.resample(time="ME").interpolate("linear") @@ -2094,7 +2111,7 @@ def test_upsample_interpolate_bug_2197(self): assert_equal(result, expected) @requires_scipy - def test_upsample_interpolate_regression_1605(self): + def test_upsample_interpolate_regression_1605(self) -> None: dates = pd.date_range("2016-01-01", "2016-03-31", freq="1D") expected = xr.DataArray( np.random.random((len(dates), 2, 3)), @@ -2107,7 +2124,7 @@ def test_upsample_interpolate_regression_1605(self): @requires_dask @requires_scipy @pytest.mark.parametrize("chunked_time", [True, False]) - def test_upsample_interpolate_dask(self, chunked_time): + def test_upsample_interpolate_dask(self, chunked_time: bool) -> None: from scipy.interpolate import interp1d xs = np.arange(6) @@ -2125,7 +2142,15 @@ def test_upsample_interpolate_dask(self, chunked_time): # Split the times into equal sub-intervals to simulate the 6 hour # to 1 hour up-sampling new_times_idx = np.linspace(0, len(times) - 1, len(times) * 5) - for kind in ["linear", "nearest", "zero", "slinear", "quadratic", "cubic"]: + kinds: list[InterpOptions] = [ + "linear", + "nearest", + "zero", + "slinear", + "quadratic", + "cubic", + ] + for kind in kinds: actual = array.chunk(chunks).resample(time="1h").interpolate(kind) actual = actual.compute() f = interp1d( @@ -2214,7 +2239,7 @@ def test_resample_invalid_loffset(self) -> None: class TestDatasetResample: - def test_resample_and_first(self): + def test_resample_and_first(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { @@ -2240,7 +2265,7 @@ def test_resample_and_first(self): result = actual.reduce(method) assert_equal(expected, result) - def test_resample_min_count(self): + def test_resample_min_count(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { @@ -2262,7 +2287,7 @@ def test_resample_min_count(self): ) assert_allclose(expected, actual) - def test_resample_by_mean_with_keep_attrs(self): + def test_resample_by_mean_with_keep_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { @@ -2282,7 +2307,7 @@ def test_resample_by_mean_with_keep_attrs(self): expected = ds.attrs assert expected == actual - def test_resample_loffset(self): + def test_resample_loffset(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { @@ -2293,7 +2318,7 @@ def test_resample_loffset(self): ) ds.attrs["dsmeta"] = "dsdata" - def test_resample_by_mean_discarding_attrs(self): + def test_resample_by_mean_discarding_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { @@ -2309,7 +2334,7 @@ def test_resample_by_mean_discarding_attrs(self): assert resampled_ds["bar"].attrs == {} assert resampled_ds.attrs == {} - def test_resample_by_last_discarding_attrs(self): + def test_resample_by_last_discarding_attrs(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { @@ -2326,7 +2351,7 @@ def test_resample_by_last_discarding_attrs(self): assert resampled_ds.attrs == {} @requires_scipy - def test_resample_drop_nondim_coords(self): + def test_resample_drop_nondim_coords(self) -> None: xs = np.arange(6) ys = np.arange(3) times = pd.date_range("2000-01-01", freq="6h", periods=5) @@ -2352,7 +2377,7 @@ def test_resample_drop_nondim_coords(self): actual = ds.resample(time="1h").interpolate("linear") assert "tc" not in actual.coords - def test_resample_old_api(self): + def test_resample_old_api(self) -> None: times = pd.date_range("2000-01-01", freq="6h", periods=10) ds = Dataset( { @@ -2363,15 +2388,15 @@ def test_resample_old_api(self): ) with pytest.raises(TypeError, match=r"resample\(\) no longer supports"): - ds.resample("1D", "time") + ds.resample("1D", "time") # type: ignore[arg-type] with pytest.raises(TypeError, match=r"resample\(\) no longer supports"): - ds.resample("1D", dim="time", how="mean") + ds.resample("1D", dim="time", how="mean") # type: ignore[arg-type] with pytest.raises(TypeError, match=r"resample\(\) no longer supports"): - ds.resample("1D", dim="time") + ds.resample("1D", dim="time") # type: ignore[arg-type] - def test_resample_ds_da_are_the_same(self): + def test_resample_ds_da_are_the_same(self) -> None: time = pd.date_range("2000-01-01", freq="6h", periods=365 * 4) ds = xr.Dataset( { @@ -2384,7 +2409,7 @@ def test_resample_ds_da_are_the_same(self): ds.resample(time="ME").mean()["foo"], ds.foo.resample(time="ME").mean() ) - def test_ds_resample_apply_func_args(self): + def test_ds_resample_apply_func_args(self) -> None: def func(arg1, arg2, arg3=0.0): return arg1.mean("time") + arg2 + arg3 @@ -2535,7 +2560,7 @@ def test_min_count_error(use_flox: bool) -> None: @requires_dask -def test_groupby_math_auto_chunk(): +def test_groupby_math_auto_chunk() -> None: da = xr.DataArray( [[1, 2, 3], [1, 2, 3], [1, 2, 3]], dims=("y", "x"), @@ -2549,7 +2574,7 @@ def test_groupby_math_auto_chunk(): @pytest.mark.parametrize("use_flox", [True, False]) -def test_groupby_dim_no_dim_equal(use_flox): +def test_groupby_dim_no_dim_equal(use_flox: bool) -> None: # https://github.com/pydata/xarray/issues/8263 da = DataArray( data=[1, 2, 3, 4], dims="lat", coords={"lat": np.linspace(0, 1.01, 4)} @@ -2561,7 +2586,7 @@ def test_groupby_dim_no_dim_equal(use_flox): @requires_flox -def test_default_flox_method(): +def test_default_flox_method() -> None: import flox.xarray da = xr.DataArray([1, 2, 3], dims="x", coords={"label": ("x", [2, 2, 1])}) diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py index 6f983a121fe..e636be5589f 100644 --- a/xarray/tests/test_plot.py +++ b/xarray/tests/test_plot.py @@ -3,9 +3,9 @@ import contextlib import inspect import math -from collections.abc import Hashable +from collections.abc import Generator, Hashable from copy import copy -from datetime import datetime +from datetime import date, datetime, timedelta from typing import Any, Callable, Literal import numpy as np @@ -85,44 +85,46 @@ def test_all_figures_closed(): @pytest.mark.flaky @pytest.mark.skip(reason="maybe flaky") -def text_in_fig(): +def text_in_fig() -> set[str]: """ Return the set of all text in the figure """ - return {t.get_text() for t in plt.gcf().findobj(mpl.text.Text)} + return {t.get_text() for t in plt.gcf().findobj(mpl.text.Text)} # type: ignore[attr-defined] # mpl error? -def find_possible_colorbars(): +def find_possible_colorbars() -> list[mpl.collections.QuadMesh]: # nb. this function also matches meshes from pcolormesh - return plt.gcf().findobj(mpl.collections.QuadMesh) + return plt.gcf().findobj(mpl.collections.QuadMesh) # type: ignore[return-value] # mpl error? -def substring_in_axes(substring, ax): +def substring_in_axes(substring: str, ax: mpl.axes.Axes) -> bool: """ Return True if a substring is found anywhere in an axes """ - alltxt = {t.get_text() for t in ax.findobj(mpl.text.Text)} + alltxt: set[str] = {t.get_text() for t in ax.findobj(mpl.text.Text)} # type: ignore[attr-defined] # mpl error? for txt in alltxt: if substring in txt: return True return False -def substring_not_in_axes(substring, ax): +def substring_not_in_axes(substring: str, ax: mpl.axes.Axes) -> bool: """ Return True if a substring is not found anywhere in an axes """ - alltxt = {t.get_text() for t in ax.findobj(mpl.text.Text)} + alltxt: set[str] = {t.get_text() for t in ax.findobj(mpl.text.Text)} # type: ignore[attr-defined] # mpl error? check = [(substring not in txt) for txt in alltxt] return all(check) -def property_in_axes_text(property, property_str, target_txt, ax): +def property_in_axes_text( + property, property_str, target_txt, ax: mpl.axes.Axes +) -> bool: """ Return True if the specified text in an axes has the property assigned to property_str """ - alltxt = ax.findobj(mpl.text.Text) + alltxt: list[mpl.text.Text] = ax.findobj(mpl.text.Text) # type: ignore[assignment] check = [] for t in alltxt: if t.get_text() == target_txt: @@ -130,7 +132,7 @@ def property_in_axes_text(property, property_str, target_txt, ax): return all(check) -def easy_array(shape, start=0, stop=1): +def easy_array(shape: tuple[int, ...], start: float = 0, stop: float = 1) -> np.ndarray: """ Make an array with desired shape using np.linspace @@ -140,7 +142,7 @@ def easy_array(shape, start=0, stop=1): return a.reshape(shape) -def get_colorbar_label(colorbar): +def get_colorbar_label(colorbar) -> str: if colorbar.orientation == "vertical": return colorbar.ax.get_ylabel() else: @@ -150,27 +152,27 @@ def get_colorbar_label(colorbar): @requires_matplotlib class PlotTestCase: @pytest.fixture(autouse=True) - def setup(self): + def setup(self) -> Generator: yield # Remove all matplotlib figures plt.close("all") - def pass_in_axis(self, plotmethod, subplot_kw=None): + def pass_in_axis(self, plotmethod, subplot_kw=None) -> None: fig, axs = plt.subplots(ncols=2, subplot_kw=subplot_kw) plotmethod(ax=axs[0]) assert axs[0].has_data() @pytest.mark.slow - def imshow_called(self, plotmethod): + def imshow_called(self, plotmethod) -> bool: plotmethod() images = plt.gca().findobj(mpl.image.AxesImage) return len(images) > 0 - def contourf_called(self, plotmethod): + def contourf_called(self, plotmethod) -> bool: plotmethod() # Compatible with mpl before (PathCollection) and after (QuadContourSet) 3.8 - def matchfunc(x): + def matchfunc(x) -> bool: return isinstance( x, (mpl.collections.PathCollection, mpl.contour.QuadContourSet) ) @@ -620,6 +622,18 @@ def test_datetime_dimension(self) -> None: ax = plt.gca() assert ax.has_data() + def test_date_dimension(self) -> None: + nrow = 3 + ncol = 4 + start = date(2000, 1, 1) + time = [start + timedelta(days=i) for i in range(nrow)] + a = DataArray( + easy_array((nrow, ncol)), coords=[("time", time), ("y", range(ncol))] + ) + a.plot() + ax = plt.gca() + assert ax.has_data() + @pytest.mark.slow @pytest.mark.filterwarnings("ignore:tight_layout cannot") def test_convenient_facetgrid(self) -> None: @@ -1236,14 +1250,16 @@ def test_discrete_colormap_list_levels_and_vmin_or_vmax(self) -> None: def test_discrete_colormap_provided_boundary_norm(self) -> None: norm = mpl.colors.BoundaryNorm([0, 5, 10, 15], 4) primitive = self.darray.plot.contourf(norm=norm) - np.testing.assert_allclose(primitive.levels, norm.boundaries) + np.testing.assert_allclose(list(primitive.levels), norm.boundaries) def test_discrete_colormap_provided_boundary_norm_matching_cmap_levels( self, ) -> None: norm = mpl.colors.BoundaryNorm([0, 5, 10, 15], 4) primitive = self.darray.plot.contourf(norm=norm) - assert primitive.colorbar.norm.Ncmap == primitive.colorbar.norm.N + cbar = primitive.colorbar + assert cbar is not None + assert cbar.norm.Ncmap == cbar.norm.N # type: ignore[attr-defined] # Exists, debatable if public though. class Common2dMixin: @@ -2028,15 +2044,17 @@ def test_normalize_rgb_one_arg_error(self) -> None: for vmin2, vmax2 in ((-1.2, -1), (2, 2.1)): da.plot.imshow(vmin=vmin2, vmax=vmax2) - def test_imshow_rgb_values_in_valid_range(self) -> None: - da = DataArray(np.arange(75, dtype="uint8").reshape((5, 5, 3))) + @pytest.mark.parametrize("dtype", [np.uint8, np.int8, np.int16]) + def test_imshow_rgb_values_in_valid_range(self, dtype) -> None: + da = DataArray(np.arange(75, dtype=dtype).reshape((5, 5, 3))) _, ax = plt.subplots() out = da.plot.imshow(ax=ax).get_array() assert out is not None - dtype = out.dtype - assert dtype is not None - assert dtype == np.uint8 + actual_dtype = out.dtype + assert actual_dtype is not None + assert actual_dtype == np.uint8 assert (out[..., :3] == da.values).all() # Compare without added alpha + assert (out[..., -1] == 255).all() # Compare alpha @pytest.mark.filterwarnings("ignore:Several dimensions of this array") def test_regression_rgb_imshow_dim_size_one(self) -> None: @@ -2518,7 +2536,7 @@ def test_default_labels(self) -> None: # Leftmost column should have array name for ax in g.axs[:, 0]: - assert substring_in_axes(self.darray.name, ax) + assert substring_in_axes(str(self.darray.name), ax) def test_test_empty_cell(self) -> None: g = ( @@ -2621,7 +2639,7 @@ def test_facetgrid(self) -> None: (True, "continuous", False, True), ], ) - def test_add_guide(self, add_guide, hue_style, legend, colorbar): + def test_add_guide(self, add_guide, hue_style, legend, colorbar) -> None: meta_data = _infer_meta_data( self.ds, x="x", @@ -2797,7 +2815,7 @@ def test_bad_args( add_legend: bool | None, add_colorbar: bool | None, error_type: type[Exception], - ): + ) -> None: with pytest.raises(error_type): self.ds.plot.scatter( x=x, y=y, hue=hue, add_legend=add_legend, add_colorbar=add_colorbar @@ -2997,20 +3015,22 @@ def test_ncaxis_notinstalled_line_plot(self) -> None: @requires_matplotlib class TestAxesKwargs: @pytest.fixture(params=[1, 2, 3]) - def data_array(self, request): + def data_array(self, request) -> DataArray: """ Return a simple DataArray """ dims = request.param if dims == 1: return DataArray(easy_array((10,))) - if dims == 2: + elif dims == 2: return DataArray(easy_array((10, 3))) - if dims == 3: + elif dims == 3: return DataArray(easy_array((10, 3, 2))) + else: + raise ValueError(f"No DataArray implemented for {dims=}.") @pytest.fixture(params=[1, 2]) - def data_array_logspaced(self, request): + def data_array_logspaced(self, request) -> DataArray: """ Return a simple DataArray with logspaced coordinates """ @@ -3019,12 +3039,14 @@ def data_array_logspaced(self, request): return DataArray( np.arange(7), dims=("x",), coords={"x": np.logspace(-3, 3, 7)} ) - if dims == 2: + elif dims == 2: return DataArray( np.arange(16).reshape(4, 4), dims=("y", "x"), coords={"x": np.logspace(-1, 2, 4), "y": np.logspace(-5, -1, 4)}, ) + else: + raise ValueError(f"No DataArray implemented for {dims=}.") @pytest.mark.parametrize("xincrease", [True, False]) def test_xincrease_kwarg(self, data_array, xincrease) -> None: @@ -3132,16 +3154,16 @@ def test_facetgrid_single_contour() -> None: @requires_matplotlib -def test_get_axis_raises(): +def test_get_axis_raises() -> None: # test get_axis raises an error if trying to do invalid things # cannot provide both ax and figsize with pytest.raises(ValueError, match="both `figsize` and `ax`"): - get_axis(figsize=[4, 4], size=None, aspect=None, ax="something") + get_axis(figsize=[4, 4], size=None, aspect=None, ax="something") # type: ignore[arg-type] # cannot provide both ax and size with pytest.raises(ValueError, match="both `size` and `ax`"): - get_axis(figsize=None, size=200, aspect=4 / 3, ax="something") + get_axis(figsize=None, size=200, aspect=4 / 3, ax="something") # type: ignore[arg-type] # cannot provide both size and figsize with pytest.raises(ValueError, match="both `figsize` and `size`"): @@ -3153,7 +3175,7 @@ def test_get_axis_raises(): # cannot provide axis and subplot_kws with pytest.raises(ValueError, match="cannot use subplot_kws with existing ax"): - get_axis(figsize=None, size=None, aspect=None, ax=1, something_else=5) + get_axis(figsize=None, size=None, aspect=None, ax=1, something_else=5) # type: ignore[arg-type] @requires_matplotlib