diff --git a/.github/ISSUE_TEMPLATE/bugreport.yml b/.github/ISSUE_TEMPLATE/bugreport.yml index 043584f3ea6..59e5889f5ec 100644 --- a/.github/ISSUE_TEMPLATE/bugreport.yml +++ b/.github/ISSUE_TEMPLATE/bugreport.yml @@ -1,4 +1,4 @@ -name: Bug Report +name: 🐛 Bug Report description: File a bug report to help us improve labels: [bug, "needs triage"] body: @@ -26,14 +26,24 @@ body: attributes: label: Minimal Complete Verifiable Example description: | - Minimal, self-contained copy-pastable example that generates the issue if possible. Please be concise with code posted. See guidelines below on how to provide a good bug report: + Minimal, self-contained copy-pastable example that demonstrates the issue. This will be automatically formatted into code, so no need for markdown backticks. + render: Python + + - type: checkboxes + id: mvce-checkboxes + attributes: + label: MVCE confirmation + description: | + Please confirm that the bug report is in an excellent state, so we can understand & fix it quickly & efficiently. For more details, check out: - [Minimal Complete Verifiable Examples](https://stackoverflow.com/help/mcve) - [Craft Minimal Bug Reports](http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports) - Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly. - This will be automatically formatted into code, so no need for markdown backticks. - render: Python + options: + - label: Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray. + - label: Complete example — the example is self-contained, including all data and the text of any traceback. + - label: Verifiable example — the example copy & pastes into an IPython prompt or [Binder notebook](https://mybinder.org/v2/gh/pydata/xarray/main?urlpath=lab/tree/doc/examples/blank_template.ipynb), returning the result. + - label: New issue — a search of GitHub Issues suggests this is not a duplicate. - type: textarea id: log-output @@ -54,6 +64,12 @@ body: attributes: label: Environment description: | - Paste the output of `xr.show_versions()` here + Paste the output of `xr.show_versions()` between the `
` tags, leaving an empty line following the opening tag. + value: | +
+ + + +
validations: required: true diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index 0ad7e5f3e13..83e5f3b97fa 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -1,8 +1,14 @@ blank_issues_enabled: false contact_links: - - name: Usage question + - name: ❓ Usage question url: https://github.com/pydata/xarray/discussions about: | Ask questions and discuss with other community members here. If you have a question like "How do I concatenate a list of datasets?" then please include a self-contained reproducible example if possible. + - name: 🗺️ Raster analysis usage question + url: https://github.com/corteva/rioxarray/discussions + about: | + If you are using the rioxarray extension (engine='rasterio'), or have questions about + raster analysis such as geospatial formats, coordinate reprojection, etc., + please use the rioxarray discussion forum. diff --git a/.github/ISSUE_TEMPLATE/misc.yml b/.github/ISSUE_TEMPLATE/misc.yml index 94dd2d86567..a98f6d90c45 100644 --- a/.github/ISSUE_TEMPLATE/misc.yml +++ b/.github/ISSUE_TEMPLATE/misc.yml @@ -1,5 +1,5 @@ -name: Issue -description: General Issue or discussion topic. For usage questions, please follow the "Usage question" link +name: 📝 Issue +description: General issue, that's not a bug report. labels: ["needs triage"] body: - type: markdown diff --git a/.github/ISSUE_TEMPLATE/newfeature.yml b/.github/ISSUE_TEMPLATE/newfeature.yml index 77cb15b7d37..04adf4bb867 100644 --- a/.github/ISSUE_TEMPLATE/newfeature.yml +++ b/.github/ISSUE_TEMPLATE/newfeature.yml @@ -1,4 +1,4 @@ -name: Feature Request +name: 💡 Feature Request description: Suggest an idea for xarray labels: [enhancement] body: diff --git a/.github/stale.yml b/.github/stale.yml index f4057844d01..e29b7ddcc5e 100644 --- a/.github/stale.yml +++ b/.github/stale.yml @@ -1,7 +1,7 @@ # Configuration for probot-stale - https://github.com/probot/stale # Number of days of inactivity before an Issue or Pull Request becomes stale -daysUntilStale: 700 # start with a large number and reduce shortly +daysUntilStale: 600 # start with a large number and reduce shortly # Number of days of inactivity before an Issue or Pull Request with the stale label is closed. # Set to false to disable. If disabled, issues still need to be closed manually, but will remain marked as stale. @@ -14,10 +14,10 @@ exemptLabels: - "[Status] Maybe Later" # Set to true to ignore issues in a project (defaults to false) -exemptProjects: false +exemptProjects: true # Set to true to ignore issues in a milestone (defaults to false) -exemptMilestones: false +exemptMilestones: true # Set to true to ignore issues with an assignee (defaults to false) exemptAssignees: true @@ -31,6 +31,9 @@ markComment: | If this issue remains relevant, please comment here or remove the `stale` label; otherwise it will be marked as closed automatically +closeComment: | + The stalebot didn't hear anything for a while, so it closed this. Please reopen if this is still an issue. + # Comment to post when removing the stale label. # unmarkComment: > # Your comment here. @@ -40,8 +43,7 @@ markComment: | # Your comment here. # Limit the number of actions per hour, from 1-30. Default is 30 -limitPerRun: 1 # start with a small number - +limitPerRun: 2 # start with a small number # Limit to only `issues` or `pulls` # only: issues diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 6d482445f96..034ffee40ad 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -67,7 +67,7 @@ jobs: cp benchmarks/README_CI.md benchmarks.log .asv/results/ working-directory: ${{ env.ASV_DIR }} - - uses: actions/upload-artifact@v2 + - uses: actions/upload-artifact@v3 if: always() with: name: asv-benchmark-results-${{ runner.os }} diff --git a/.github/workflows/ci-additional.yaml b/.github/workflows/ci-additional.yaml index 50c95cdebb7..ff3e8ab7e63 100644 --- a/.github/workflows/ci-additional.yaml +++ b/.github/workflows/ci-additional.yaml @@ -30,51 +30,19 @@ jobs: with: keyword: "[skip-ci]" - test: - name: ${{ matrix.os }} ${{ matrix.env }} - runs-on: ${{ matrix.os }} + doctest: + name: Doctests + runs-on: "ubuntu-latest" needs: detect-ci-trigger if: needs.detect-ci-trigger.outputs.triggered == 'false' defaults: run: shell: bash -l {0} - strategy: - fail-fast: false - matrix: - os: ["ubuntu-latest"] - env: - [ - # Minimum python version: - "py38-bare-minimum", - "py38-min-all-deps", - # Latest python version: - "py39-all-but-dask", - "py39-flaky", - ] steps: - uses: actions/checkout@v3 with: fetch-depth: 0 # Fetch all history for all branches and tags. - - - name: Set environment variables - run: | - if [[ ${{ matrix.env }} == "py39-flaky" ]] ; - then - echo "CONDA_ENV_FILE=ci/requirements/environment.yml" >> $GITHUB_ENV - echo "PYTEST_EXTRA_FLAGS=--run-flaky --run-network-tests" >> $GITHUB_ENV - - else - echo "CONDA_ENV_FILE=ci/requirements/${{ matrix.env }}.yml" >> $GITHUB_ENV - fi - - name: Cache conda - uses: actions/cache@v2 - with: - path: ~/conda_pkgs_dir - key: - ${{ runner.os }}-conda-${{ matrix.env }}-${{ - hashFiles('ci/requirements/**.yml') }} - - uses: conda-incubator/setup-miniconda@v2 with: channels: conda-forge @@ -82,43 +50,28 @@ jobs: mamba-version: "*" activate-environment: xarray-tests auto-update-conda: false - python-version: 3.9 - use-only-tar-bz2: true + python-version: "3.9" - name: Install conda dependencies run: | - mamba env update -f $CONDA_ENV_FILE - + mamba env update -f ci/requirements/environment.yml - 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 - - name: Import xarray - run: | - python -c "import xarray" - - name: Run tests + - name: Run doctests run: | - python -m pytest -n 4 \ - --cov=xarray \ - --cov-report=xml \ - $PYTEST_EXTRA_FLAGS + python -m pytest --doctest-modules xarray --ignore xarray/tests - - name: Upload code coverage to Codecov - uses: codecov/codecov-action@v2.1.0 - with: - file: ./coverage.xml - flags: unittests,${{ matrix.env }} - env_vars: RUNNER_OS - name: codecov-umbrella - fail_ci_if_error: false - doctest: - name: Doctests + mypy: + name: Mypy runs-on: "ubuntu-latest" + needs: detect-ci-trigger + # temporarily skipping due to https://github.com/pydata/xarray/issues/6551 if: needs.detect-ci-trigger.outputs.triggered == 'false' defaults: run: @@ -148,9 +101,13 @@ jobs: conda info -a conda list python xarray/util/print_versions.py - - name: Run doctests + - name: Install mypy run: | - python -m pytest --doctest-modules xarray --ignore xarray/tests + python -m pip install mypy + + - name: Run mypy + run: | + python -m mypy --install-types --non-interactive min-version-policy: name: Minimum Version Policy @@ -176,5 +133,5 @@ jobs: - name: minimum versions policy run: | mamba install -y pyyaml conda python-dateutil - python ci/min_deps_check.py ci/requirements/py38-bare-minimum.yml - python ci/min_deps_check.py ci/requirements/py38-min-all-deps.yml + python ci/min_deps_check.py ci/requirements/bare-minimum.yml + python ci/min_deps_check.py ci/requirements/min-all-deps.yml diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 4747b5ae20d..20f876b52fc 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -30,7 +30,7 @@ jobs: with: keyword: "[skip-ci]" test: - name: ${{ matrix.os }} py${{ matrix.python-version }} + name: ${{ matrix.os }} py${{ matrix.python-version }} ${{ matrix.env }} runs-on: ${{ matrix.os }} needs: detect-ci-trigger if: needs.detect-ci-trigger.outputs.triggered == 'false' @@ -42,7 +42,23 @@ jobs: matrix: os: ["ubuntu-latest", "macos-latest", "windows-latest"] # Bookend python versions - python-version: ["3.8", "3.9", "3.10"] + python-version: ["3.8", "3.10"] + env: [""] + include: + # Minimum python version: + - env: "bare-minimum" + python-version: "3.8" + os: ubuntu-latest + - env: "min-all-deps" + python-version: "3.8" + os: ubuntu-latest + # Latest python version: + - env: "all-but-dask" + python-version: "3.10" + os: ubuntu-latest + - env: "flaky" + python-version: "3.10" + os: ubuntu-latest steps: - uses: actions/checkout@v3 with: @@ -52,18 +68,27 @@ jobs: if [[ ${{ matrix.os }} == windows* ]] ; then echo "CONDA_ENV_FILE=ci/requirements/environment-windows.yml" >> $GITHUB_ENV + elif [[ "${{ matrix.env }}" != "" ]] ; + then + if [[ "${{ matrix.env }}" == "flaky" ]] ; + then + echo "CONDA_ENV_FILE=ci/requirements/environment.yml" >> $GITHUB_ENV + echo "PYTEST_EXTRA_FLAGS=--run-flaky --run-network-tests" >> $GITHUB_ENV + else + echo "CONDA_ENV_FILE=ci/requirements/${{ matrix.env }}.yml" >> $GITHUB_ENV + fi else echo "CONDA_ENV_FILE=ci/requirements/environment.yml" >> $GITHUB_ENV - fi + echo "PYTHON_VERSION=${{ matrix.python-version }}" >> $GITHUB_ENV - name: Cache conda - uses: actions/cache@v2 + uses: actions/cache@v3 with: path: ~/conda_pkgs_dir - key: ${{ runner.os }}-conda-py${{ matrix.python-version }}-${{ - hashFiles('ci/requirements/**.yml') }} + key: ${{ runner.os }}-conda-py${{ matrix.python-version }}-${{ hashFiles('ci/requirements/**.yml') }}-${{ matrix.env }} + - uses: conda-incubator/setup-miniconda@v2 with: channels: conda-forge @@ -78,6 +103,13 @@ jobs: run: | mamba env update -f $CONDA_ENV_FILE + # We only want to install this on one run, because otherwise we'll have + # duplicate annotations. + - name: Install error reporter + if: ${{ matrix.os }} == 'ubuntu-latest' and ${{ matrix.python-version }} == '3.10' + run: | + python -m pip install pytest-github-actions-annotate-failures + - name: Install xarray run: | python -m pip install --no-deps -e . @@ -87,24 +119,27 @@ jobs: conda info -a conda list python xarray/util/print_versions.py + - name: Import xarray run: | python -c "import xarray" + - name: Run tests run: python -m pytest -n 4 --cov=xarray --cov-report=xml --junitxml=pytest.xml + $PYTEST_EXTRA_FLAGS - name: Upload test results if: always() - uses: actions/upload-artifact@v2 + uses: actions/upload-artifact@v3 with: name: Test results for ${{ runner.os }}-${{ matrix.python-version }} path: pytest.xml - name: Upload code coverage to Codecov - uses: codecov/codecov-action@v2.1.0 + uses: codecov/codecov-action@v3.1.0 with: file: ./coverage.xml flags: unittests @@ -118,7 +153,7 @@ jobs: if: github.repository == 'pydata/xarray' steps: - name: Upload - uses: actions/upload-artifact@v2 + uses: actions/upload-artifact@v3 with: name: Event File path: ${{ github.event_path }} diff --git a/.github/workflows/pypi-release.yaml b/.github/workflows/pypi-release.yaml index c88cf556a50..9cad271ce6f 100644 --- a/.github/workflows/pypi-release.yaml +++ b/.github/workflows/pypi-release.yaml @@ -41,7 +41,7 @@ jobs: else echo "✅ Looks good" fi - - uses: actions/upload-artifact@v2 + - uses: actions/upload-artifact@v3 with: name: releases path: dist @@ -54,7 +54,7 @@ jobs: name: Install Python with: python-version: 3.8 - - uses: actions/download-artifact@v2 + - uses: actions/download-artifact@v3 with: name: releases path: dist @@ -85,7 +85,7 @@ jobs: if: github.event_name == 'release' runs-on: ubuntu-latest steps: - - uses: actions/download-artifact@v2 + - uses: actions/download-artifact@v3 with: name: releases path: dist diff --git a/.github/workflows/upstream-dev-ci.yaml b/.github/workflows/upstream-dev-ci.yaml index 6091306ed8b..81d1c7db4b8 100644 --- a/.github/workflows/upstream-dev-ci.yaml +++ b/.github/workflows/upstream-dev-ci.yaml @@ -92,7 +92,7 @@ jobs: && steps.status.outcome == 'failure' && github.event_name == 'schedule' && github.repository == 'pydata/xarray' - uses: actions/upload-artifact@v2 + uses: actions/upload-artifact@v3 with: name: output-${{ matrix.python-version }}-log path: output-${{ matrix.python-version }}-log @@ -114,7 +114,7 @@ jobs: - uses: actions/setup-python@v3 with: python-version: "3.x" - - uses: actions/download-artifact@v2 + - uses: actions/download-artifact@v3 with: path: /tmp/workspace/logs - name: Move all log files into a single directory diff --git a/.gitignore b/.gitignore index 686c7efa701..21c18c17ff7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ *.py[cod] __pycache__ +.env +.venv # example caches from Hypothesis .hypothesis/ @@ -73,3 +75,7 @@ xarray/tests/data/*.grib.*.idx Icon* .ipynb_checkpoints +doc/team-panel.txt +doc/external-examples-gallery.txt +doc/notebooks-examples-gallery.txt +doc/videos-gallery.txt diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e92479ded86..427af33b833 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,7 +1,7 @@ # https://pre-commit.com/ repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.1.0 + rev: v4.2.0 hooks: - id: trailing-whitespace - id: end-of-file-fixer @@ -19,14 +19,14 @@ repos: hooks: - id: isort - repo: https://github.com/asottile/pyupgrade - rev: v2.31.1 + rev: v2.32.1 hooks: - id: pyupgrade args: - "--py38-plus" # https://github.com/python/black#version-control-integration - repo: https://github.com/psf/black - rev: 22.1.0 + rev: 22.3.0 hooks: - id: black - id: black-jupyter @@ -35,6 +35,7 @@ repos: hooks: - id: blackdoc exclude: "generate_reductions.py" + additional_dependencies: ["black==22.3.0"] - repo: https://github.com/PyCQA/flake8 rev: 4.0.1 hooks: @@ -45,7 +46,7 @@ repos: # - id: velin # args: ["--write", "--compact"] - repo: https://github.com/pre-commit/mirrors-mypy - rev: v0.942 + rev: v0.960 hooks: - id: mypy # Copied from setup.cfg @@ -62,3 +63,7 @@ repos: typing-extensions==3.10.0.0, numpy, ] + - repo: https://github.com/citation-file-format/cff-converter-python + rev: ebf0b5e44d67f8beaa1cd13a0d0393ea04c6058d + hooks: + - id: validate-cff diff --git a/HOW_TO_RELEASE.md b/HOW_TO_RELEASE.md index 8d82277ae55..f647263a3a7 100644 --- a/HOW_TO_RELEASE.md +++ b/HOW_TO_RELEASE.md @@ -111,4 +111,4 @@ upstream https://github.com/pydata/xarray (push) As of 2022.03.0, we utilize the [CALVER](https://calver.org/) version system. Specifically, we have adopted the pattern `YYYY.MM.X`, where `YYYY` is a 4-digit -year (e.g. `2022`), `MM` is a 2-digit zero-padded month (e.g. `01` for January), and `X` is the release number (starting at zero at the start of each month and incremented once for each additional release). +year (e.g. `2022`), `0M` is a 2-digit zero-padded month (e.g. `01` for January), and `X` is the release number (starting at zero at the start of each month and incremented once for each additional release). diff --git a/README.md b/README.md new file mode 100644 index 00000000000..4c7306214bb --- /dev/null +++ b/README.md @@ -0,0 +1,138 @@ +# xarray: N-D labeled arrays and datasets + +[![CI](https://github.com/pydata/xarray/workflows/CI/badge.svg?branch=main)](https://github.com/pydata/xarray/actions?query=workflow%3ACI) +[![Code coverage](https://codecov.io/gh/pydata/xarray/branch/main/graph/badge.svg)](https://codecov.io/gh/pydata/xarray) +[![Docs](https://readthedocs.org/projects/xray/badge/?version=latest)](https://docs.xarray.dev/) +[![Benchmarked with asv](https://img.shields.io/badge/benchmarked%20by-asv-green.svg?style=flat)](https://pandas.pydata.org/speed/xarray/) +[![Available on pypi](https://img.shields.io/pypi/v/xarray.svg)](https://pypi.python.org/pypi/xarray/) +[![Formatted with black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/python/black) +[![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/) +[![Mirror on zendoo](https://zenodo.org/badge/DOI/10.5281/zenodo.598201.svg)](https://doi.org/10.5281/zenodo.598201) +[![Examples on binder](https://img.shields.io/badge/launch-binder-579ACA.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAFkAAABZCAMAAABi1XidAAAB8lBMVEX///9XmsrmZYH1olJXmsr1olJXmsrmZYH1olJXmsr1olJXmsrmZYH1olL1olJXmsr1olJXmsrmZYH1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olJXmsrmZYH1olL1olL0nFf1olJXmsrmZYH1olJXmsq8dZb1olJXmsrmZYH1olJXmspXmspXmsr1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olLeaIVXmsrmZYH1olL1olL1olJXmsrmZYH1olLna31Xmsr1olJXmsr1olJXmsrmZYH1olLqoVr1olJXmsr1olJXmsrmZYH1olL1olKkfaPobXvviGabgadXmsqThKuofKHmZ4Dobnr1olJXmsr1olJXmspXmsr1olJXmsrfZ4TuhWn1olL1olJXmsqBi7X1olJXmspZmslbmMhbmsdemsVfl8ZgmsNim8Jpk8F0m7R4m7F5nLB6jbh7jbiDirOEibOGnKaMhq+PnaCVg6qWg6qegKaff6WhnpKofKGtnomxeZy3noG6dZi+n3vCcpPDcpPGn3bLb4/Mb47UbIrVa4rYoGjdaIbeaIXhoWHmZYHobXvpcHjqdHXreHLroVrsfG/uhGnuh2bwj2Hxk17yl1vzmljzm1j0nlX1olL3AJXWAAAAbXRSTlMAEBAQHx8gICAuLjAwMDw9PUBAQEpQUFBXV1hgYGBkcHBwcXl8gICAgoiIkJCQlJicnJ2goKCmqK+wsLC4usDAwMjP0NDQ1NbW3Nzg4ODi5+3v8PDw8/T09PX29vb39/f5+fr7+/z8/Pz9/v7+zczCxgAABC5JREFUeAHN1ul3k0UUBvCb1CTVpmpaitAGSLSpSuKCLWpbTKNJFGlcSMAFF63iUmRccNG6gLbuxkXU66JAUef/9LSpmXnyLr3T5AO/rzl5zj137p136BISy44fKJXuGN/d19PUfYeO67Znqtf2KH33Id1psXoFdW30sPZ1sMvs2D060AHqws4FHeJojLZqnw53cmfvg+XR8mC0OEjuxrXEkX5ydeVJLVIlV0e10PXk5k7dYeHu7Cj1j+49uKg7uLU61tGLw1lq27ugQYlclHC4bgv7VQ+TAyj5Zc/UjsPvs1sd5cWryWObtvWT2EPa4rtnWW3JkpjggEpbOsPr7F7EyNewtpBIslA7p43HCsnwooXTEc3UmPmCNn5lrqTJxy6nRmcavGZVt/3Da2pD5NHvsOHJCrdc1G2r3DITpU7yic7w/7Rxnjc0kt5GC4djiv2Sz3Fb2iEZg41/ddsFDoyuYrIkmFehz0HR2thPgQqMyQYb2OtB0WxsZ3BeG3+wpRb1vzl2UYBog8FfGhttFKjtAclnZYrRo9ryG9uG/FZQU4AEg8ZE9LjGMzTmqKXPLnlWVnIlQQTvxJf8ip7VgjZjyVPrjw1te5otM7RmP7xm+sK2Gv9I8Gi++BRbEkR9EBw8zRUcKxwp73xkaLiqQb+kGduJTNHG72zcW9LoJgqQxpP3/Tj//c3yB0tqzaml05/+orHLksVO+95kX7/7qgJvnjlrfr2Ggsyx0eoy9uPzN5SPd86aXggOsEKW2Prz7du3VID3/tzs/sSRs2w7ovVHKtjrX2pd7ZMlTxAYfBAL9jiDwfLkq55Tm7ifhMlTGPyCAs7RFRhn47JnlcB9RM5T97ASuZXIcVNuUDIndpDbdsfrqsOppeXl5Y+XVKdjFCTh+zGaVuj0d9zy05PPK3QzBamxdwtTCrzyg/2Rvf2EstUjordGwa/kx9mSJLr8mLLtCW8HHGJc2R5hS219IiF6PnTusOqcMl57gm0Z8kanKMAQg0qSyuZfn7zItsbGyO9QlnxY0eCuD1XL2ys/MsrQhltE7Ug0uFOzufJFE2PxBo/YAx8XPPdDwWN0MrDRYIZF0mSMKCNHgaIVFoBbNoLJ7tEQDKxGF0kcLQimojCZopv0OkNOyWCCg9XMVAi7ARJzQdM2QUh0gmBozjc3Skg6dSBRqDGYSUOu66Zg+I2fNZs/M3/f/Grl/XnyF1Gw3VKCez0PN5IUfFLqvgUN4C0qNqYs5YhPL+aVZYDE4IpUk57oSFnJm4FyCqqOE0jhY2SMyLFoo56zyo6becOS5UVDdj7Vih0zp+tcMhwRpBeLyqtIjlJKAIZSbI8SGSF3k0pA3mR5tHuwPFoa7N7reoq2bqCsAk1HqCu5uvI1n6JuRXI+S1Mco54YmYTwcn6Aeic+kssXi8XpXC4V3t7/ADuTNKaQJdScAAAAAElFTkSuQmCC)](https://mybinder.org/v2/gh/pydata/xarray/main?urlpath=lab/tree/doc/examples/weather-data.ipynb) +[![Twitter](https://img.shields.io/twitter/follow/xarray_dev?style=social)](https://twitter.com/xarray_dev) + +**xarray** (formerly **xray**) is an open source project and Python +package that makes working with labelled multi-dimensional arrays +simple, efficient, and fun! + +Xarray introduces labels in the form of dimensions, coordinates and +attributes on top of raw [NumPy](https://www.numpy.org)-like arrays, +which allows for a more intuitive, more concise, and less error-prone +developer experience. The package includes a large and growing library +of domain-agnostic functions for advanced analytics and visualization +with these data structures. + +Xarray was inspired by and borrows heavily from +[pandas](https://pandas.pydata.org), the popular data analysis package +focused on labelled tabular data. It is particularly tailored to working +with [netCDF](https://www.unidata.ucar.edu/software/netcdf) files, which +were the source of xarray\'s data model, and integrates tightly with +[dask](https://dask.org) for parallel computing. + +## Why xarray? + +Multi-dimensional (a.k.a. N-dimensional, ND) arrays (sometimes called +"tensors") are an essential part of computational science. They are +encountered in a wide range of fields, including physics, astronomy, +geoscience, bioinformatics, engineering, finance, and deep learning. In +Python, [NumPy](https://www.numpy.org) provides the fundamental data +structure and API for working with raw ND arrays. However, real-world +datasets are usually more than just raw numbers; they have labels which +encode information about how the array values map to locations in space, +time, etc. + +Xarray doesn\'t just keep track of labels on arrays \-- it uses them to +provide a powerful and concise interface. For example: + +- Apply operations over dimensions by name: `x.sum('time')`. +- Select values by label instead of integer location: + `x.loc['2014-01-01']` or `x.sel(time='2014-01-01')`. +- Mathematical operations (e.g., `x - y`) vectorize across multiple + dimensions (array broadcasting) based on dimension names, not shape. +- Flexible split-apply-combine operations with groupby: + `x.groupby('time.dayofyear').mean()`. +- Database like alignment based on coordinate labels that smoothly + handles missing values: `x, y = xr.align(x, y, join='outer')`. +- Keep track of arbitrary metadata in the form of a Python dictionary: + `x.attrs`. + +## Documentation + +Learn more about xarray in its official documentation at +. + +Try out an [interactive Jupyter +notebook](https://mybinder.org/v2/gh/pydata/xarray/main?urlpath=lab/tree/doc/examples/weather-data.ipynb). + +## Contributing + +You can find information about contributing to xarray at our +[Contributing +page](https://docs.xarray.dev/en/latest/contributing.html#). + +## Get in touch + +- Ask usage questions ("How do I?") on + [GitHub Discussions](https://github.com/pydata/xarray/discussions). +- Report bugs, suggest features or view the source code [on + GitHub](https://github.com/pydata/xarray). +- For less well defined questions or ideas, or to announce other + projects of interest to xarray users, use the [mailing + list](https://groups.google.com/forum/#!forum/xarray). + +## NumFOCUS + + + +Xarray is a fiscally sponsored project of +[NumFOCUS](https://numfocus.org), a nonprofit dedicated to supporting +the open source scientific computing community. If you like Xarray and +want to support our mission, please consider making a +[donation](https://numfocus.salsalabs.org/donate-to-xarray/) to support +our efforts. + +## History + +Xarray is an evolution of an internal tool developed at [The Climate +Corporation](http://climate.com/). It was originally written by Climate +Corp researchers Stephan Hoyer, Alex Kleeman and Eugene Brevdo and was +released as open source in May 2014. The project was renamed from +"xray" in January 2016. Xarray became a fiscally sponsored project of +[NumFOCUS](https://numfocus.org) in August 2018. + +## Contributors + +Thanks to our many contributors! + +[![Contributors](https://contrib.rocks/image?repo=pydata/xarray)](https://github.com/pydata/xarray/graphs/contributors) + +## License + +Copyright 2014-2019, xarray Developers + +Licensed under the Apache License, Version 2.0 (the "License"); you +may not use this file except in compliance with the License. You may +obtain a copy of the License at + + + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Xarray bundles portions of pandas, NumPy and Seaborn, all of which are +available under a "3-clause BSD" license: + +- pandas: setup.py, xarray/util/print_versions.py +- NumPy: xarray/core/npcompat.py +- Seaborn: _determine_cmap_params in xarray/core/plot/utils.py + +Xarray also bundles portions of CPython, which is available under the +"Python Software Foundation License" in xarray/core/pycompat.py. + +Xarray uses icons from the icomoon package (free version), which is +available under the "CC BY 4.0" license. + +The full text of these licenses are included in the licenses directory. diff --git a/README.rst b/README.rst deleted file mode 100644 index 7a4ad4e1f9f..00000000000 --- a/README.rst +++ /dev/null @@ -1,146 +0,0 @@ -xarray: N-D labeled arrays and datasets -======================================= - -.. image:: https://github.com/pydata/xarray/workflows/CI/badge.svg?branch=main - :target: https://github.com/pydata/xarray/actions?query=workflow%3ACI -.. image:: https://codecov.io/gh/pydata/xarray/branch/main/graph/badge.svg - :target: https://codecov.io/gh/pydata/xarray -.. image:: https://readthedocs.org/projects/xray/badge/?version=latest - :target: https://docs.xarray.dev/ -.. image:: https://img.shields.io/badge/benchmarked%20by-asv-green.svg?style=flat - :target: https://pandas.pydata.org/speed/xarray/ -.. image:: https://img.shields.io/pypi/v/xarray.svg - :target: https://pypi.python.org/pypi/xarray/ -.. image:: https://img.shields.io/badge/code%20style-black-000000.svg - :target: https://github.com/python/black -.. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.598201.svg - :target: https://doi.org/10.5281/zenodo.598201 - - -**xarray** (formerly **xray**) is an open source project and Python package -that makes working with labelled multi-dimensional arrays simple, -efficient, and fun! - -Xarray introduces labels in the form of dimensions, coordinates and -attributes on top of raw NumPy_-like arrays, which allows for a more -intuitive, more concise, and less error-prone developer experience. -The package includes a large and growing library of domain-agnostic functions -for advanced analytics and visualization with these data structures. - -Xarray was inspired by and borrows heavily from pandas_, the popular data -analysis package focused on labelled tabular data. -It is particularly tailored to working with netCDF_ files, which were the -source of xarray's data model, and integrates tightly with dask_ for parallel -computing. - -.. _NumPy: https://www.numpy.org -.. _pandas: https://pandas.pydata.org -.. _dask: https://dask.org -.. _netCDF: https://www.unidata.ucar.edu/software/netcdf - -Why xarray? ------------ - -Multi-dimensional (a.k.a. N-dimensional, ND) arrays (sometimes called -"tensors") are an essential part of computational science. -They are encountered in a wide range of fields, including physics, astronomy, -geoscience, bioinformatics, engineering, finance, and deep learning. -In Python, NumPy_ provides the fundamental data structure and API for -working with raw ND arrays. -However, real-world datasets are usually more than just raw numbers; -they have labels which encode information about how the array values map -to locations in space, time, etc. - -Xarray doesn't just keep track of labels on arrays -- it uses them to provide a -powerful and concise interface. For example: - -- Apply operations over dimensions by name: ``x.sum('time')``. -- Select values by label instead of integer location: - ``x.loc['2014-01-01']`` or ``x.sel(time='2014-01-01')``. -- Mathematical operations (e.g., ``x - y``) vectorize across multiple - dimensions (array broadcasting) based on dimension names, not shape. -- Flexible split-apply-combine operations with groupby: - ``x.groupby('time.dayofyear').mean()``. -- Database like alignment based on coordinate labels that smoothly - handles missing values: ``x, y = xr.align(x, y, join='outer')``. -- Keep track of arbitrary metadata in the form of a Python dictionary: - ``x.attrs``. - -Documentation -------------- - -Learn more about xarray in its official documentation at https://docs.xarray.dev/ - -Contributing ------------- - -You can find information about contributing to xarray at our `Contributing page `_. - -Get in touch ------------- - -- Ask usage questions ("How do I?") on `StackOverflow`_. -- Report bugs, suggest features or view the source code `on GitHub`_. -- For less well defined questions or ideas, or to announce other projects of - interest to xarray users, use the `mailing list`_. - -.. _StackOverFlow: https://stackoverflow.com/questions/tagged/python-xarray -.. _mailing list: https://groups.google.com/forum/#!forum/xarray -.. _on GitHub: https://github.com/pydata/xarray - -NumFOCUS --------- - -.. image:: https://numfocus.org/wp-content/uploads/2017/07/NumFocus_LRG.png - :scale: 25 % - :target: https://numfocus.org/ - -Xarray is a fiscally sponsored project of NumFOCUS_, a nonprofit dedicated -to supporting the open source scientific computing community. If you like -Xarray and want to support our mission, please consider making a donation_ -to support our efforts. - -.. _donation: https://numfocus.salsalabs.org/donate-to-xarray/ - -History -------- - -Xarray is an evolution of an internal tool developed at `The Climate -Corporation`__. It was originally written by Climate Corp researchers Stephan -Hoyer, Alex Kleeman and Eugene Brevdo and was released as open source in -May 2014. The project was renamed from "xray" in January 2016. Xarray became a -fiscally sponsored project of NumFOCUS_ in August 2018. - -__ http://climate.com/ -.. _NumFOCUS: https://numfocus.org - -License -------- - -Copyright 2014-2019, xarray Developers - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - https://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - -Xarray bundles portions of pandas, NumPy and Seaborn, all of which are available -under a "3-clause BSD" license: -- pandas: setup.py, xarray/util/print_versions.py -- NumPy: xarray/core/npcompat.py -- Seaborn: _determine_cmap_params in xarray/core/plot/utils.py - -Xarray also bundles portions of CPython, which is available under the "Python -Software Foundation License" in xarray/core/pycompat.py. - -Xarray uses icons from the icomoon package (free version), which is -available under the "CC BY 4.0" license. - -The full text of these licenses are included in the licenses directory. diff --git a/asv_bench/asv.conf.json b/asv_bench/asv.conf.json index 10b8aead374..5de6d6a4f76 100644 --- a/asv_bench/asv.conf.json +++ b/asv_bench/asv.conf.json @@ -58,6 +58,8 @@ // "pip+emcee": [""], // emcee is only available for install with pip. // }, "matrix": { + "setuptools_scm[toml]": [""], // GH6609 + "setuptools_scm_git_archive": [""], // GH6609 "numpy": [""], "pandas": [""], "netcdf4": [""], diff --git a/asv_bench/benchmarks/polyfit.py b/asv_bench/benchmarks/polyfit.py new file mode 100644 index 00000000000..429ffa19baa --- /dev/null +++ b/asv_bench/benchmarks/polyfit.py @@ -0,0 +1,38 @@ +import numpy as np + +import xarray as xr + +from . import parameterized, randn, requires_dask + +NDEGS = (2, 5, 20) +NX = (10**2, 10**6) + + +class Polyval: + def setup(self, *args, **kwargs): + self.xs = {nx: xr.DataArray(randn((nx,)), dims="x", name="x") for nx in NX} + self.coeffs = { + ndeg: xr.DataArray( + randn((ndeg,)), dims="degree", coords={"degree": np.arange(ndeg)} + ) + for ndeg in NDEGS + } + + @parameterized(["nx", "ndeg"], [NX, NDEGS]) + def time_polyval(self, nx, ndeg): + x = self.xs[nx] + c = self.coeffs[ndeg] + xr.polyval(x, c).compute() + + @parameterized(["nx", "ndeg"], [NX, NDEGS]) + def peakmem_polyval(self, nx, ndeg): + x = self.xs[nx] + c = self.coeffs[ndeg] + xr.polyval(x, c).compute() + + +class PolyvalDask(Polyval): + def setup(self, *args, **kwargs): + requires_dask() + super().setup(*args, **kwargs) + self.xs = {k: v.chunk({"x": 10000}) for k, v in self.xs.items()} diff --git a/ci/requirements/py39-all-but-dask.yml b/ci/requirements/all-but-dask.yml similarity index 90% rename from ci/requirements/py39-all-but-dask.yml rename to ci/requirements/all-but-dask.yml index f05745ee1fa..2381fdcdfe0 100644 --- a/ci/requirements/py39-all-but-dask.yml +++ b/ci/requirements/all-but-dask.yml @@ -3,7 +3,7 @@ channels: - conda-forge - nodefaults dependencies: - - python=3.9 + - python=3.10 - black - aiobotocore - boto3 @@ -18,7 +18,7 @@ dependencies: - h5py - hdf5 - hypothesis - - lxml # Optional dep of pydap + - lxml # Optional dep of pydap - matplotlib-base - nc-time-axis - netcdf4 @@ -44,4 +44,4 @@ dependencies: - typing_extensions - zarr - pip: - - numbagg + - numbagg diff --git a/ci/requirements/py38-bare-minimum.yml b/ci/requirements/bare-minimum.yml similarity index 100% rename from ci/requirements/py38-bare-minimum.yml rename to ci/requirements/bare-minimum.yml diff --git a/ci/requirements/doc.yml b/ci/requirements/doc.yml index 3d4e149619e..437c493c92c 100644 --- a/ci/requirements/doc.yml +++ b/ci/requirements/doc.yml @@ -13,7 +13,6 @@ dependencies: - ipykernel - ipython - iris>=2.3 - - jinja2<3.1 # remove once nbconvert fixed the use of removed functions - jupyter_client - matplotlib-base - nbsphinx @@ -33,7 +32,7 @@ dependencies: - sphinx-autosummary-accessors - sphinx-book-theme >= 0.0.38 - sphinx-copybutton - - sphinx-panels + - sphinx-design - sphinx!=4.4.0 - zarr>=2.4 - pip: diff --git a/ci/requirements/environment-windows.yml b/ci/requirements/environment-windows.yml index a31188fec5b..44dfcc173fb 100644 --- a/ci/requirements/environment-windows.yml +++ b/ci/requirements/environment-windows.yml @@ -10,6 +10,7 @@ dependencies: - cftime - dask-core - distributed + - flox - fsspec!=2021.7.0 - h5netcdf - h5py diff --git a/ci/requirements/environment.yml b/ci/requirements/environment.yml index 1ac4f87d34a..7e62576cfdc 100644 --- a/ci/requirements/environment.yml +++ b/ci/requirements/environment.yml @@ -12,6 +12,7 @@ dependencies: - cftime - dask-core - distributed + - flox - fsspec!=2021.7.0 - h5netcdf - h5py @@ -38,7 +39,6 @@ dependencies: - pytest - pytest-cov - pytest-env - - pytest-github-actions-annotate-failures - pytest-xdist - rasterio - scipy diff --git a/ci/requirements/py38-min-all-deps.yml b/ci/requirements/min-all-deps.yml similarity index 72% rename from ci/requirements/py38-min-all-deps.yml rename to ci/requirements/min-all-deps.yml index 76e2b28093d..34879af730b 100644 --- a/ci/requirements/py38-min-all-deps.yml +++ b/ci/requirements/min-all-deps.yml @@ -10,46 +10,46 @@ dependencies: - python=3.8 - boto3=1.13 - bottleneck=1.3 - # cartopy 0.18 conflicts with pynio - - cartopy=0.17 + - cartopy=0.19 - cdms2=3.1 - cfgrib=0.9 - - cftime=1.2 + - cftime=1.4 - coveralls - - dask-core=2.30 - - distributed=2.30 - - h5netcdf=0.8 - - h5py=2.10 - # hdf5 1.12 conflicts with h5py=2.10 + - dask-core=2021.04 + - distributed=2021.04 + - flox=0.5 + - h5netcdf=0.11 + - h5py=3.1 + # hdf5 1.12 conflicts with h5py=3.1 - hdf5=1.10 - hypothesis - iris=2.4 - lxml=4.6 # Optional dep of pydap - - matplotlib-base=3.3 + - matplotlib-base=3.4 - nc-time-axis=1.2 # netcdf follows a 1.major.minor[.patch] convention # (see https://github.com/Unidata/netcdf4-python/issues/1090) # bumping the netCDF4 version is currently blocked by #4491 - netcdf4=1.5.3 - - numba=0.51 - - numpy=1.18 + - numba=0.53 + - numpy=1.19 - packaging=20.0 - - pandas=1.1 - - pint=0.16 + - pandas=1.2 + - pint=0.17 - pip - pseudonetcdf=3.1 - pydap=3.2 - - pynio=1.5 + # - pynio=1.5.5 - pytest - pytest-cov - pytest-env - pytest-xdist - - rasterio=1.1 - - scipy=1.5 + - rasterio=1.2 + - scipy=1.6 - seaborn=0.11 - - sparse=0.11 + - sparse=0.12 - toolz=0.11 - typing_extensions=3.7 - - zarr=2.5 + - zarr=2.8 - pip: - numbagg==0.1 diff --git a/doc/_static/index_api.svg b/doc/_static/index_api.svg new file mode 100644 index 00000000000..87013d24ce3 --- /dev/null +++ b/doc/_static/index_api.svg @@ -0,0 +1,97 @@ + + + + + + + + + + image/svg+xml + + + + + + + + + + + + + + + + + diff --git a/doc/_static/index_contribute.svg b/doc/_static/index_contribute.svg new file mode 100644 index 00000000000..399f1d7630b --- /dev/null +++ b/doc/_static/index_contribute.svg @@ -0,0 +1,76 @@ + + + + + + + + + + image/svg+xml + + + + + + + + + + + + diff --git a/doc/_static/index_getting_started.svg b/doc/_static/index_getting_started.svg new file mode 100644 index 00000000000..d1c7b08a2a1 --- /dev/null +++ b/doc/_static/index_getting_started.svg @@ -0,0 +1,66 @@ + + + + + + + + + + image/svg+xml + + + + + + + + + diff --git a/doc/_static/index_user_guide.svg b/doc/_static/index_user_guide.svg new file mode 100644 index 00000000000..bff2482423d --- /dev/null +++ b/doc/_static/index_user_guide.svg @@ -0,0 +1,67 @@ + + + + + + + + + + image/svg+xml + + + + + + + + + diff --git a/doc/_static/style.css b/doc/_static/style.css index 833b11a83ab..3c21d7ac7c9 100644 --- a/doc/_static/style.css +++ b/doc/_static/style.css @@ -261,3 +261,13 @@ body { .bd-toc .nav > .active > ul { display: block; } + +/* Main index page overview cards */ + +.sd-card-img-top { + width: 33% !important; + display: block; + margin-left: auto; + margin-right: auto; + margin-top: 10px; +} diff --git a/doc/api-hidden.rst b/doc/api-hidden.rst index 8ed9e47be01..30bc9f858f2 100644 --- a/doc/api-hidden.rst +++ b/doc/api-hidden.rst @@ -319,66 +319,6 @@ IndexVariable.sizes IndexVariable.values - ufuncs.angle - ufuncs.arccos - ufuncs.arccosh - ufuncs.arcsin - ufuncs.arcsinh - ufuncs.arctan - ufuncs.arctan2 - ufuncs.arctanh - ufuncs.ceil - ufuncs.conj - ufuncs.copysign - ufuncs.cos - ufuncs.cosh - ufuncs.deg2rad - ufuncs.degrees - ufuncs.exp - ufuncs.expm1 - ufuncs.fabs - ufuncs.fix - ufuncs.floor - ufuncs.fmax - ufuncs.fmin - ufuncs.fmod - ufuncs.fmod - ufuncs.frexp - ufuncs.hypot - ufuncs.imag - ufuncs.iscomplex - ufuncs.isfinite - ufuncs.isinf - ufuncs.isnan - ufuncs.isreal - ufuncs.ldexp - ufuncs.log - ufuncs.log10 - ufuncs.log1p - ufuncs.log2 - ufuncs.logaddexp - ufuncs.logaddexp2 - ufuncs.logical_and - ufuncs.logical_not - ufuncs.logical_or - ufuncs.logical_xor - ufuncs.maximum - ufuncs.minimum - ufuncs.nextafter - ufuncs.rad2deg - ufuncs.radians - ufuncs.real - ufuncs.rint - ufuncs.sign - ufuncs.signbit - ufuncs.sin - ufuncs.sinh - ufuncs.sqrt - ufuncs.square - ufuncs.tan - ufuncs.tanh - ufuncs.trunc - plot.plot plot.line plot.step diff --git a/doc/api.rst b/doc/api.rst index 7fdd775e168..644b86cdebb 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -610,84 +610,6 @@ Plotting DataArray.plot.step DataArray.plot.surface -.. _api.ufuncs: - -Universal functions -=================== - -.. warning:: - - With recent versions of NumPy, Dask and xarray, NumPy ufuncs are now - supported directly on all xarray and Dask objects. This obviates the need - for the ``xarray.ufuncs`` module, which should not be used for new code - unless compatibility with versions of NumPy prior to v1.13 is - required. They will be removed once support for NumPy prior to - v1.17 is dropped. - -These functions are copied from NumPy, but extended to work on NumPy arrays, -dask arrays and all xarray objects. You can find them in the ``xarray.ufuncs`` -module: - -:py:attr:`~ufuncs.angle` -:py:attr:`~ufuncs.arccos` -:py:attr:`~ufuncs.arccosh` -:py:attr:`~ufuncs.arcsin` -:py:attr:`~ufuncs.arcsinh` -:py:attr:`~ufuncs.arctan` -:py:attr:`~ufuncs.arctan2` -:py:attr:`~ufuncs.arctanh` -:py:attr:`~ufuncs.ceil` -:py:attr:`~ufuncs.conj` -:py:attr:`~ufuncs.copysign` -:py:attr:`~ufuncs.cos` -:py:attr:`~ufuncs.cosh` -:py:attr:`~ufuncs.deg2rad` -:py:attr:`~ufuncs.degrees` -:py:attr:`~ufuncs.exp` -:py:attr:`~ufuncs.expm1` -:py:attr:`~ufuncs.fabs` -:py:attr:`~ufuncs.fix` -:py:attr:`~ufuncs.floor` -:py:attr:`~ufuncs.fmax` -:py:attr:`~ufuncs.fmin` -:py:attr:`~ufuncs.fmod` -:py:attr:`~ufuncs.fmod` -:py:attr:`~ufuncs.frexp` -:py:attr:`~ufuncs.hypot` -:py:attr:`~ufuncs.imag` -:py:attr:`~ufuncs.iscomplex` -:py:attr:`~ufuncs.isfinite` -:py:attr:`~ufuncs.isinf` -:py:attr:`~ufuncs.isnan` -:py:attr:`~ufuncs.isreal` -:py:attr:`~ufuncs.ldexp` -:py:attr:`~ufuncs.log` -:py:attr:`~ufuncs.log10` -:py:attr:`~ufuncs.log1p` -:py:attr:`~ufuncs.log2` -:py:attr:`~ufuncs.logaddexp` -:py:attr:`~ufuncs.logaddexp2` -:py:attr:`~ufuncs.logical_and` -:py:attr:`~ufuncs.logical_not` -:py:attr:`~ufuncs.logical_or` -:py:attr:`~ufuncs.logical_xor` -:py:attr:`~ufuncs.maximum` -:py:attr:`~ufuncs.minimum` -:py:attr:`~ufuncs.nextafter` -:py:attr:`~ufuncs.rad2deg` -:py:attr:`~ufuncs.radians` -:py:attr:`~ufuncs.real` -:py:attr:`~ufuncs.rint` -:py:attr:`~ufuncs.sign` -:py:attr:`~ufuncs.signbit` -:py:attr:`~ufuncs.sin` -:py:attr:`~ufuncs.sinh` -:py:attr:`~ufuncs.sqrt` -:py:attr:`~ufuncs.square` -:py:attr:`~ufuncs.tan` -:py:attr:`~ufuncs.tanh` -:py:attr:`~ufuncs.trunc` - IO / Conversion =============== diff --git a/doc/conf.py b/doc/conf.py index 8ce9efdce88..7e28953bc7f 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -15,14 +15,21 @@ import datetime import inspect import os +import pathlib import subprocess import sys from contextlib import suppress +from textwrap import dedent, indent import sphinx_autosummary_accessors +import yaml +from sphinx.application import Sphinx +from sphinx.util import logging import xarray +LOGGER = logging.getLogger("conf") + allowed_failures = set() print("python exec:", sys.executable) @@ -60,6 +67,8 @@ ] ) +nbsphinx_allow_errors = True + # -- General configuration ------------------------------------------------ # If your documentation needs a minimal Sphinx version, state it here. @@ -80,12 +89,13 @@ "nbsphinx", "sphinx_autosummary_accessors", "sphinx.ext.linkcode", - "sphinx_panels", "sphinxext.opengraph", "sphinx_copybutton", "sphinxext.rediraffe", + "sphinx_design", ] + extlinks = { "issue": ("https://github.com/pydata/xarray/issues/%s", "GH"), "pull": ("https://github.com/pydata/xarray/pull/%s", "PR"), @@ -186,7 +196,7 @@ # General information about the project. project = "xarray" -copyright = "2014-%s, xarray Developers" % datetime.datetime.now().year +copyright = f"2014-{datetime.datetime.now().year}, xarray Developers" # The short X.Y version. version = xarray.__version__.split("+")[0] @@ -383,5 +393,109 @@ def html_page_context(app, pagename, templatename, context, doctree): context["theme_use_edit_page_button"] = False -def setup(app): +def update_team(app: Sphinx): + """Update the team members list.""" + + LOGGER.info("Updating team members page...") + + team = yaml.safe_load(pathlib.Path(app.srcdir, "team.yml").read_bytes()) + items = [] + for member in team: + item = f""" + .. grid-item-card:: + :text-align: center + :link: https://github.com/{member['gh_login']} + + .. image:: {member['avatar']} + :alt: {member['name']} + +++ + {member['name']} + """ + items.append(item) + + items_md = indent(dedent("\n".join(items)), prefix=" ") + + markdown = f""" +.. grid:: 1 2 3 3 + :gutter: 2 + + {items_md} + """ + + pathlib.Path(app.srcdir, "team-panel.txt").write_text(markdown) + LOGGER.info("Team members page updated.") + + +def update_gallery(app: Sphinx): + """Update the gallery page.""" + + LOGGER.info("Updating gallery page...") + + gallery = yaml.safe_load(pathlib.Path(app.srcdir, "gallery.yml").read_bytes()) + + for key in gallery: + items = [ + f""" + .. grid-item-card:: + :text-align: center + :link: {item['path']} + + .. image:: {item['thumbnail']} + :alt: {item['title']} + +++ + {item['title']} + """ + for item in gallery[key] + ] + + items_md = indent(dedent("\n".join(items)), prefix=" ") + markdown = f""" +.. grid:: 1 2 2 2 + :gutter: 2 + + {items_md} + """ + pathlib.Path(app.srcdir, f"{key}-gallery.txt").write_text(markdown) + LOGGER.info(f"{key} gallery page updated.") + LOGGER.info("Gallery page updated.") + + +def update_videos(app: Sphinx): + """Update the videos page.""" + + LOGGER.info("Updating videos page...") + + videos = yaml.safe_load(pathlib.Path(app.srcdir, "videos.yml").read_bytes()) + + items = [] + for video in videos: + + authors = " | ".join(video["authors"]) + item = f""" +.. grid-item-card:: {" ".join(video["title"].split())} + :text-align: center + + .. raw:: html + + {video['src']} + +++ + {authors} + """ + items.append(item) + + items_md = indent(dedent("\n".join(items)), prefix=" ") + markdown = f""" +.. grid:: 1 2 2 2 + :gutter: 2 + + {items_md} + """ + pathlib.Path(app.srcdir, "videos-gallery.txt").write_text(markdown) + LOGGER.info("Videos page updated.") + + +def setup(app: Sphinx): app.connect("html-page-context", html_page_context) + app.connect("builder-inited", update_team) + app.connect("builder-inited", update_gallery) + app.connect("builder-inited", update_videos) diff --git a/doc/ecosystem.rst b/doc/ecosystem.rst index 2b49b1529e1..61b60ab9e83 100644 --- a/doc/ecosystem.rst +++ b/doc/ecosystem.rst @@ -74,6 +74,7 @@ Extend xarray capabilities - `nxarray `_: NeXus input/output capability for xarray. - `xarray-compare `_: xarray extension for data comparison. - `xarray-dataclasses `_: xarray extension for typed DataArray and Dataset creation. +- `xarray_einstats `_: Statistics, linear algebra and einops for xarray - `xarray_extras `_: Advanced algorithms for xarray objects (e.g. integrations/interpolations). - `xpublish `_: Publish Xarray Datasets via a Zarr compatible REST API. - `xrft `_: Fourier transforms for xarray data. diff --git a/doc/examples/blank_template.ipynb b/doc/examples/blank_template.ipynb new file mode 100644 index 00000000000..bcb15c1158d --- /dev/null +++ b/doc/examples/blank_template.ipynb @@ -0,0 +1,58 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d8f54f6a", + "metadata": {}, + "source": [ + "# Blank template\n", + "\n", + "Use this notebook from Binder to test an issue or reproduce a bug report" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "41b90ede", + "metadata": {}, + "outputs": [], + "source": [ + "import xarray as xr\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "ds = xr.tutorial.load_dataset(\"air_temperature\")\n", + "da = ds[\"air\"]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "effd9aeb", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.10" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/doc/gallery.rst b/doc/gallery.rst index 36eb39d1a53..61ec45c7bda 100644 --- a/doc/gallery.rst +++ b/doc/gallery.rst @@ -10,75 +10,7 @@ Contributions are highly welcomed and appreciated. So, if you are interested in Notebook Examples ----------------- -.. panels:: - :column: text-center col-lg-6 col-md-6 col-sm-12 col-xs-12 p-2 - :card: +my-2 - :img-top-cls: w-75 m-auto p-2 - :body: d-none - - --- - :img-top: _static/thumbnails/toy-weather-data.png - ++++ - .. link-button:: examples/weather-data - :type: ref - :text: Toy weather data - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/thumbnails/monthly-means.png - ++++ - .. link-button:: examples/monthly-means - :type: ref - :text: Calculating Seasonal Averages from Timeseries of Monthly Means - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/thumbnails/area_weighted_temperature.png - ++++ - .. link-button:: examples/area_weighted_temperature - :type: ref - :text: Compare weighted and unweighted mean temperature - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/thumbnails/multidimensional-coords.png - ++++ - .. link-button:: examples/multidimensional-coords - :type: ref - :text: Working with Multidimensional Coordinates - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/thumbnails/visualization_gallery.png - ++++ - .. link-button:: examples/visualization_gallery - :type: ref - :text: Visualization Gallery - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/thumbnails/ROMS_ocean_model.png - ++++ - .. link-button:: examples/ROMS_ocean_model - :type: ref - :text: ROMS Ocean Model Example - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/thumbnails/ERA5-GRIB-example.png - ++++ - .. link-button:: examples/ERA5-GRIB-example - :type: ref - :text: GRIB Data Example - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/dataset-diagram-square-logo.png - ++++ - .. link-button:: examples/apply_ufunc_vectorize_1d - :type: ref - :text: Applying unvectorized functions with apply_ufunc - :classes: btn-outline-dark btn-block stretched-link +.. include:: notebooks-examples-gallery.txt .. toctree:: @@ -93,38 +25,11 @@ Notebook Examples examples/ROMS_ocean_model examples/ERA5-GRIB-example examples/apply_ufunc_vectorize_1d + examples/blank_template External Examples ----------------- -.. panels:: - :column: text-center col-lg-6 col-md-6 col-sm-12 col-xs-12 p-2 - :card: +my-2 - :img-top-cls: w-75 m-auto p-2 - :body: d-none - - --- - :img-top: _static/dataset-diagram-square-logo.png - ++++ - .. link-button:: https://corteva.github.io/rioxarray/stable/examples/examples.html - :type: url - :text: Managing raster data with rioxarray - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: https://avatars.githubusercontent.com/u/60833341?s=200&v=4 - ++++ - .. link-button:: https://gallery.pangeo.io/ - :type: url - :text: Xarray and dask on the cloud with Pangeo - :classes: btn-outline-dark btn-block stretched-link - - --- - :img-top: _static/dataset-diagram-square-logo.png - ++++ - .. link-button:: https://examples.dask.org/xarray.html - :type: url - :text: Xarray with Dask Arrays - :classes: btn-outline-dark btn-block stretched-link +.. include:: external-examples-gallery.txt diff --git a/doc/gallery.yml b/doc/gallery.yml new file mode 100644 index 00000000000..f1a147dae87 --- /dev/null +++ b/doc/gallery.yml @@ -0,0 +1,45 @@ +notebooks-examples: + - title: Toy weather data + path: examples/weather-data.html + thumbnail: _static/thumbnails/toy-weather-data.png + + - title: Calculating Seasonal Averages from Timeseries of Monthly Means + path: examples/monthly-means.html + thumbnail: _static/thumbnails/monthly-means.png + + - title: Compare weighted and unweighted mean temperature + path: examples/area_weighted_temperature.html + thumbnail: _static/thumbnails/area_weighted_temperature.png + + - title: Working with Multidimensional Coordinates + path: examples/multidimensional-coords.html + thumbnail: _static/thumbnails/multidimensional-coords.png + + - title: Visualization Gallery + path: examples/visualization_gallery.html + thumbnail: _static/thumbnails/visualization_gallery.png + + - title: GRIB Data Example + path: examples/ERA5-GRIB-example.html + thumbnail: _static/thumbnails/ERA5-GRIB-example.png + + - title: Applying unvectorized functions with apply_ufunc + path: examples/apply_ufunc_vectorize_1d.html + thumbnail: _static/dataset-diagram-square-logo.png + +external-examples: + - title: Managing raster data with rioxarray + path: https://corteva.github.io/rioxarray/stable/examples/examples.html + thumbnail: _static/dataset-diagram-square-logo.png + + - title: Xarray and dask on the cloud with Pangeo + path: https://gallery.pangeo.io/ + thumbnail: https://avatars.githubusercontent.com/u/60833341?s=200&v=4 + + - title: Xarray with Dask Arrays + path: https://examples.dask.org/xarray.html_ + thumbnail: _static/dataset-diagram-square-logo.png + + - title: Project Pythia Foundations Book + path: https://foundations.projectpythia.org/core/xarray.html + thumbnail: https://raw.githubusercontent.com/ProjectPythia/projectpythia.github.io/main/portal/_static/images/logos/pythia_logo-blue-btext-twocolor.svg diff --git a/doc/getting-started-guide/installing.rst b/doc/getting-started-guide/installing.rst index 6177ba0aaac..faa0fba5dd3 100644 --- a/doc/getting-started-guide/installing.rst +++ b/doc/getting-started-guide/installing.rst @@ -7,9 +7,9 @@ Required dependencies --------------------- - Python (3.8 or later) -- `numpy `__ (1.18 or later) +- `numpy `__ (1.19 or later) - `packaging `__ (20.0 or later) -- `pandas `__ (1.1 or later) +- `pandas `__ (1.2 or later) .. _optional-dependencies: @@ -102,7 +102,7 @@ release is guaranteed to work. You can see the actual minimum tested versions: -``_ +``_ .. _installation-instructions: diff --git a/doc/index.rst b/doc/index.rst index c549c33aa62..973f4c2c6d1 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -1,27 +1,56 @@ -xarray: N-D labeled arrays and datasets in Python -================================================= +Xarray documentation +==================== -**xarray** (formerly **xray**) is an open source project and Python package -that makes working with labelled multi-dimensional arrays simple, +Xarray makes working with labelled multi-dimensional arrays in Python simple, efficient, and fun! -Xarray introduces labels in the form of dimensions, coordinates and -attributes on top of raw NumPy_-like arrays, which allows for a more -intuitive, more concise, and less error-prone developer experience. -The package includes a large and growing library of domain-agnostic functions -for advanced analytics and visualization with these data structures. +**Useful links**: +`Home `__ | +`Code Repository `__ | +`Issues `__ | +`Discussions `__ | +`Releases `__ | +`Stack Overflow `__ | +`Mailing List `__ | +`Blog `__ -Xarray is inspired by and borrows heavily from pandas_, the popular data -analysis package focused on labelled tabular data. -It is particularly tailored to working with netCDF_ files, which were the -source of xarray's data model, and integrates tightly with dask_ for parallel -computing. -.. _NumPy: https://www.numpy.org -.. _pandas: https://pandas.pydata.org -.. _dask: https://dask.org -.. _netCDF: https://www.unidata.ucar.edu/software/netcdf +.. grid:: 1 1 2 2 + :gutter: 2 + .. grid-item-card:: Getting started + :img-top: _static/index_getting_started.svg + :link: getting-started-guide/index + :link-type: doc + + New to *xarray*? Check out the getting started guides. They contain an + introduction to *Xarray's* main concepts and links to additional tutorials. + + .. grid-item-card:: User guide + :img-top: _static/index_user_guide.svg + :link: user-guide/index + :link-type: doc + + The user guide provides in-depth information on the + key concepts of Xarray with useful background information and explanation. + + .. grid-item-card:: API reference + :img-top: _static/index_api.svg + :link: api + :link-type: doc + + The reference guide contains a detailed description of the Xarray API. + The reference describes how the methods work and which parameters can + be used. It assumes that you have an understanding of the key concepts. + + .. grid-item-card:: Developer guide + :img-top: _static/index_contribute.svg + :link: contributing + :link-type: doc + + Saw a typo in the documentation? Want to improve existing functionalities? + The contributing guidelines will guide you through the process of improving + Xarray. .. toctree:: :maxdepth: 2 @@ -56,54 +85,3 @@ computing. GitHub discussions StackOverflow - - - - -Get in touch ------------- - -- If you have a question like "How do I concatenate a list of datasets?", ask on `GitHub discussions`_ or `StackOverflow`_. - Please include a self-contained reproducible example if possible. -- Report bugs, suggest features or view the source code `on GitHub`_. -- For less well defined questions or ideas, or to announce other projects of - interest to xarray users, use `GitHub discussions`_ or the `mailing list`_. - -.. _StackOverFlow: https://stackoverflow.com/questions/tagged/python-xarray -.. _Github discussions: https://github.com/pydata/xarray/discussions -.. _mailing list: https://groups.google.com/forum/#!forum/xarray -.. _on GitHub: https://github.com/pydata/xarray - -NumFOCUS --------- - -.. image:: _static/numfocus_logo.png - :scale: 50 % - :target: https://numfocus.org/ - -Xarray is a fiscally sponsored project of NumFOCUS_, a nonprofit dedicated -to supporting the open source scientific computing community. If you like -Xarray and want to support our mission, please consider making a donation_ -to support our efforts. - -.. _donation: https://numfocus.salsalabs.org/donate-to-xarray/ - - -History -------- - -Xarray is an evolution of an internal tool developed at `The Climate -Corporation`__. It was originally written by Climate Corp researchers Stephan -Hoyer, Alex Kleeman and Eugene Brevdo and was released as open source in -May 2014. The project was renamed from "xray" in January 2016. Xarray became a -fiscally sponsored project of NumFOCUS_ in August 2018. - -__ https://climate.com/ -.. _NumFOCUS: https://numfocus.org - -License -------- - -Xarray is available under the open source `Apache License`__. - -__ https://www.apache.org/licenses/LICENSE-2.0.html diff --git a/doc/internals/extending-xarray.rst b/doc/internals/extending-xarray.rst index 2951ce10f21..f8b61d12a2f 100644 --- a/doc/internals/extending-xarray.rst +++ b/doc/internals/extending-xarray.rst @@ -92,8 +92,8 @@ on ways to write new accessors and the philosophy behind the approach, see To help users keep things straight, please `let us know `_ if you plan to write a new accessor -for an open source library. In the future, we will maintain a list of accessors -and the libraries that implement them on this page. +for an open source library. Existing open source accessors and the libraries +that implement them are available in the list on the :ref:`ecosystem` page. To make documenting accessors with ``sphinx`` and ``sphinx.ext.autosummary`` easier, you can use `sphinx-autosummary-accessors`_. diff --git a/doc/internals/how-to-add-new-backend.rst b/doc/internals/how-to-add-new-backend.rst index 506a8eb21be..bb497a1c062 100644 --- a/doc/internals/how-to-add-new-backend.rst +++ b/doc/internals/how-to-add-new-backend.rst @@ -439,27 +439,25 @@ currently available in :py:mod:`~xarray.backends` module. .. _RST preferred_chunks: -Backend preferred chunks -^^^^^^^^^^^^^^^^^^^^^^^^ - -The backend is not directly involved in `Dask `__ -chunking, since it is internally managed by Xarray. However, the backend can -define the preferred chunk size inside the variable’s encoding -``var.encoding["preferred_chunks"]``. The ``preferred_chunks`` may be useful -to improve performances with lazy loading. ``preferred_chunks`` shall be a -dictionary specifying chunk size per dimension like -``{“dim1”: 1000, “dim2”: 2000}`` or -``{“dim1”: [1000, 100], “dim2”: [2000, 2000, 2000]]}``. - -The ``preferred_chunks`` is used by Xarray to define the chunk size in some -special cases: - -- if ``chunks`` along a dimension is ``None`` or not defined -- if ``chunks`` is ``"auto"``. - -In the first case Xarray uses the chunks size specified in -``preferred_chunks``. -In the second case Xarray accommodates ideal chunk sizes, preserving if -possible the "preferred_chunks". The ideal chunk size is computed using -:py:func:`dask.array.core.normalize_chunks`, setting -``previous_chunks = preferred_chunks``. +Preferred chunk sizes +^^^^^^^^^^^^^^^^^^^^^ + +To potentially improve performance with lazy loading, the backend may define for each +variable the chunk sizes that it prefers---that is, sizes that align with how the +variable is stored. (Note that the backend is not directly involved in `Dask +`__ chunking, because Xarray internally manages chunking.) To define +the preferred chunk sizes, store a mapping within the variable's encoding under the key +``"preferred_chunks"`` (that is, ``var.encoding["preferred_chunks"]``). The mapping's +keys shall be the names of dimensions with preferred chunk sizes, and each value shall +be the corresponding dimension's preferred chunk sizes expressed as either an integer +(such as ``{"dim1": 1000, "dim2": 2000}``) or a tuple of integers (such as ``{"dim1": +(1000, 100), "dim2": (2000, 2000, 2000)}``). + +Xarray uses the preferred chunk sizes in some special cases of the ``chunks`` argument +of the :py:func:`~xarray.open_dataset` and :py:func:`~xarray.open_mfdataset` functions. +If ``chunks`` is a ``dict``, then for any dimensions missing from the keys or whose +value is ``None``, Xarray sets the chunk sizes to the preferred sizes. If ``chunks`` +equals ``"auto"``, then Xarray seeks ideal chunk sizes informed by the preferred chunk +sizes. Specifically, it determines the chunk sizes using +:py:func:`dask.array.core.normalize_chunks` with the ``previous_chunks`` argument set +according to the preferred chunk sizes. diff --git a/doc/internals/zarr-encoding-spec.rst b/doc/internals/zarr-encoding-spec.rst index f8bffa6e82f..7f468b8b0db 100644 --- a/doc/internals/zarr-encoding-spec.rst +++ b/doc/internals/zarr-encoding-spec.rst @@ -32,9 +32,11 @@ the variable dimension names and then removed from the attributes dictionary returned to the user. Because of these choices, Xarray cannot read arbitrary array data, but only -Zarr data with valid ``_ARRAY_DIMENSIONS`` attributes on each array. +Zarr data with valid ``_ARRAY_DIMENSIONS`` or +`NCZarr `_ attributes +on each array (NCZarr dimension names are defined in the ``.zarray`` file). -After decoding the ``_ARRAY_DIMENSIONS`` attribute and assigning the variable +After decoding the ``_ARRAY_DIMENSIONS`` or NCZarr attribute and assigning the variable dimensions, Xarray proceeds to [optionally] decode each variable using its standard CF decoding machinery used for NetCDF data (see :py:func:`decode_cf`). diff --git a/doc/team.rst b/doc/team.rst index 937d16627c7..ce0cb33c6ba 100644 --- a/doc/team.rst +++ b/doc/team.rst @@ -8,84 +8,7 @@ Xarray core developers are responsible for the ongoing organizational maintenanc The current core developers team comprises: -.. panels:: - :column: col-lg-4 col-md-4 col-sm-6 col-xs-12 p-2 - :card: text-center - - --- - .. image:: https://avatars.githubusercontent.com/u/1217238?v=4 - +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/shoyer,"Stephan Hoyer",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/1197350?v=4 - +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/rabernat,"Ryan Abernathey",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/2443309?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/jhamman,"Joe Hamman",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/4160723?v=4 - +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/benbovy,"Benoit Bovy",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/10050469?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/fmaussion,"Fabien Maussion",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/6815844?v=4 - +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/fujiisoup,"Keisuke Fujii",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/5635139?v=4 - +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/max-sixty,"Maximilian Roos",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/2448579?v=4 - +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/dcherian,"Deepak Cherian",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/6628425?v=4 - +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/spencerkclark,"Spencer Clark",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/35968931?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/TomNicholas,"Tom Nicholas",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/6213168?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/crusaderky,"Guido Imperiale",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/14808389?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/keewis,"Justus Magin",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/10194086?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/mathause,"Mathias Hauser",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/13301940?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/andersy005,"Anderson Banihirwe",cls=btn badge-light` - - --- - .. image:: https://avatars.githubusercontent.com/u/14371165?v=4 - ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - :link-badge:`https://github.com/Illviljan,"Jimmy Westling",cls=btn badge-light` +.. include:: team-panel.txt The full list of contributors is on our `GitHub Contributors Page `__. diff --git a/doc/team.yml b/doc/team.yml new file mode 100644 index 00000000000..56ccc457420 --- /dev/null +++ b/doc/team.yml @@ -0,0 +1,59 @@ +- name: Stephan Hoyer + gh_login: shoyer + avatar: https://avatars.githubusercontent.com/u/1217238?v=4 + +- name: Joe Hamman + gh_login: jhamman + avatar: https://avatars.githubusercontent.com/u/2443309?v=4 + +- name: Ryan Abernathey + gh_login: rabernat + avatar: https://avatars.githubusercontent.com/u/1197350?v=4 + +- name: Benoit Bovy + gh_login: genbovy + avatar: https://avatars.githubusercontent.com/u/4160723?v=4 + +- name: Fabien Maussion + gh_login: fmaussion + avatar: https://avatars.githubusercontent.com/u/10050469?v=4 + +- name: Keisuke Fujii + gh_login: fujiisoup + avatar: https://avatars.githubusercontent.com/u/6815844?v=4 + +- name: Maximilian Roos + gh_login: max-sixty + avatar: https://avatars.githubusercontent.com/u/5635139?v=4 + +- name: Deepak Cherian + gh_login: dcherian + avatar: https://avatars.githubusercontent.com/u/2448579?v=4 + +- name: Spencer Clark + gh_login: spencerkclark + avatar: https://avatars.githubusercontent.com/u/6628425?v=4 + +- name: Tom Nicholas + gh_login: TomNicholas + avatar: https://avatars.githubusercontent.com/u/35968931?v=4 + +- name: Guido Imperiale + gh_login: crusaderky + avatar: https://avatars.githubusercontent.com/u/6213168?v=4 + +- name: Justus Magin + gh_login: keewis + avatar: https://avatars.githubusercontent.com/u/14808389?v=4 + +- name: Mathias Hauser + gh_login: mathause + avatar: https://avatars.githubusercontent.com/u/10194086?v=4 + +- name: Anderson Banihirwe + gh_login: andersy005 + avatar: https://avatars.githubusercontent.com/u/13301940?v=4 + +- name: Jimmy Westling + gh_login: Illviljan + avatar: https://avatars.githubusercontent.com/u/14371165?v=4 diff --git a/doc/tutorials-and-videos.rst b/doc/tutorials-and-videos.rst index 37b8d1968eb..7a9e524340b 100644 --- a/doc/tutorials-and-videos.rst +++ b/doc/tutorials-and-videos.rst @@ -15,59 +15,8 @@ Tutorials Videos ------- -.. panels:: - :card: text-center +.. include:: videos-gallery.txt - --- - Xdev Python Tutorial Seminar Series 2022 Thinking with Xarray : High-level computation patterns | Deepak Cherian - ^^^ - .. raw:: html - - - - --- - Xdev Python Tutorial Seminar Series 2021 seminar introducing xarray (2 of 2) | Anderson Banihirwe - ^^^ - .. raw:: html - - - - --- - Xdev Python Tutorial Seminar Series 2021 seminar introducing xarray (1 of 2) | Anderson Banihirwe - ^^^ - .. raw:: html - - - - --- - Xarray's virtual tutorial | October 2020 | Anderson Banihirwe, Deepak Cherian, and Martin Durant - ^^^ - .. raw:: html - - - - --- - Xarray's Tutorial presented at the 2020 SciPy Conference | Joe Hamman, Ryan Abernathey, - Deepak Cherian, and Stephan Hoyer - ^^^ - .. raw:: html - - - - --- - Scipy 2015 talk introducing xarray to a general audience | Stephan Hoyer - ^^^ - .. raw:: html - - - - --- - 2015 Unidata Users Workshop talk and tutorial with (`with answers`_) introducing - xarray to users familiar with netCDF | Stephan Hoyer - ^^^ - .. raw:: html - - Books, Chapters and Articles ----------------------------- diff --git a/doc/user-guide/combining.rst b/doc/user-guide/combining.rst index 06dc11cea10..1dad2009665 100644 --- a/doc/user-guide/combining.rst +++ b/doc/user-guide/combining.rst @@ -22,10 +22,10 @@ Combining data Concatenate ~~~~~~~~~~~ -To combine arrays along existing or new dimension into a larger array, you -can use :py:func:`~xarray.concat`. ``concat`` takes an iterable of ``DataArray`` -or ``Dataset`` objects, as well as a dimension name, and concatenates along -that dimension: +To combine :py:class:`~xarray.Dataset`s / :py:class:`~xarray.DataArray`s along an existing or new dimension +into a larger object, you can use :py:func:`~xarray.concat`. ``concat`` +takes an iterable of ``DataArray`` or ``Dataset`` objects, as well as a +dimension name, and concatenates along that dimension: .. ipython:: python diff --git a/doc/user-guide/dask.rst b/doc/user-guide/dask.rst index 5110a970390..56717f5306e 100644 --- a/doc/user-guide/dask.rst +++ b/doc/user-guide/dask.rst @@ -84,7 +84,7 @@ argument to :py:func:`~xarray.open_dataset` or using the In this example ``latitude`` and ``longitude`` do not appear in the ``chunks`` dict, so only one chunk will be used along those dimensions. It is also -entirely equivalent to opening a dataset using :py:meth:`~xarray.open_dataset` +entirely equivalent to opening a dataset using :py:func:`~xarray.open_dataset` and then chunking the data using the ``chunk`` method, e.g., ``xr.open_dataset('example-data.nc').chunk({'time': 10})``. @@ -95,13 +95,21 @@ use :py:func:`~xarray.open_mfdataset`:: This function will automatically concatenate and merge datasets into one in the simple cases that it understands (see :py:func:`~xarray.combine_by_coords` -for the full disclaimer). By default, :py:meth:`~xarray.open_mfdataset` will chunk each +for the full disclaimer). By default, :py:func:`~xarray.open_mfdataset` will chunk each netCDF file into a single Dask array; again, supply the ``chunks`` argument to control the size of the resulting Dask arrays. In more complex cases, you can -open each file individually using :py:meth:`~xarray.open_dataset` and merge the result, as -described in :ref:`combining data`. Passing the keyword argument ``parallel=True`` to :py:meth:`~xarray.open_mfdataset` will speed up the reading of large multi-file datasets by +open each file individually using :py:func:`~xarray.open_dataset` and merge the result, as +described in :ref:`combining data`. Passing the keyword argument ``parallel=True`` to +:py:func:`~xarray.open_mfdataset` will speed up the reading of large multi-file datasets by executing those read tasks in parallel using ``dask.delayed``. +.. warning:: + + :py:func:`~xarray.open_mfdataset` called without ``chunks`` argument will return + dask arrays with chunk sizes equal to the individual files. Re-chunking + the dataset after creation with ``ds.chunk()`` will lead to an ineffective use of + memory and is not recommended. + You'll notice that printing a dataset still shows a preview of array values, even if they are actually Dask arrays. We can do this quickly with Dask because we only need to compute the first few values (typically from the first block). @@ -224,6 +232,7 @@ disk. available memory. .. note:: + For more on the differences between :py:meth:`~xarray.Dataset.persist` and :py:meth:`~xarray.Dataset.compute` see this `Stack Overflow answer `_ and the `Dask documentation `_. @@ -236,6 +245,11 @@ sizes of Dask arrays is done with the :py:meth:`~xarray.Dataset.chunk` method: rechunked = ds.chunk({"latitude": 100, "longitude": 100}) +.. warning:: + + Rechunking an existing dask array created with :py:func:`~xarray.open_mfdataset` + is not recommended (see above). + You can view the size of existing chunks on an array by viewing the :py:attr:`~xarray.Dataset.chunks` attribute: @@ -295,8 +309,7 @@ each block of your xarray object, you have three options: ``apply_ufunc`` ~~~~~~~~~~~~~~~ -Another option is to use xarray's :py:func:`~xarray.apply_ufunc`, which can -automate `embarrassingly parallel +:py:func:`~xarray.apply_ufunc` automates `embarrassingly parallel `__ "map" type operations where a function written for processing NumPy arrays should be repeatedly applied to xarray objects containing Dask arrays. It works similarly to @@ -542,18 +555,20 @@ larger chunksizes. Optimization Tips ----------------- -With analysis pipelines involving both spatial subsetting and temporal resampling, Dask performance can become very slow in certain cases. Here are some optimization tips we have found through experience: +With analysis pipelines involving both spatial subsetting and temporal resampling, Dask performance +can become very slow or memory hungry in certain cases. Here are some optimization tips we have found +through experience: -1. Do your spatial and temporal indexing (e.g. ``.sel()`` or ``.isel()``) early in the pipeline, especially before calling ``resample()`` or ``groupby()``. Grouping and resampling triggers some computation on all the blocks, which in theory should commute with indexing, but this optimization hasn't been implemented in Dask yet. (See `Dask issue #746 `_). +1. Do your spatial and temporal indexing (e.g. ``.sel()`` or ``.isel()``) early in the pipeline, especially before calling ``resample()`` or ``groupby()``. Grouping and resampling triggers some computation on all the blocks, which in theory should commute with indexing, but this optimization hasn't been implemented in Dask yet. (See `Dask issue #746 `_). More generally, ``groupby()`` is a costly operation and does not (yet) perform well on datasets split across multiple files (see :pull:`5734` and linked discussions there). 2. Save intermediate results to disk as a netCDF files (using ``to_netcdf()``) and then load them again with ``open_dataset()`` for further computations. For example, if subtracting temporal mean from a dataset, save the temporal mean to disk before subtracting. Again, in theory, Dask should be able to do the computation in a streaming fashion, but in practice this is a fail case for the Dask scheduler, because it tries to keep every chunk of an array that it computes in memory. (See `Dask issue #874 `_) -3. Specify smaller chunks across space when using :py:meth:`~xarray.open_mfdataset` (e.g., ``chunks={'latitude': 10, 'longitude': 10}``). This makes spatial subsetting easier, because there's no risk you will load chunks of data referring to different chunks (probably not necessary if you follow suggestion 1). +3. Specify smaller chunks across space when using :py:meth:`~xarray.open_mfdataset` (e.g., ``chunks={'latitude': 10, 'longitude': 10}``). This makes spatial subsetting easier, because there's no risk you will load subsets of data which span multiple chunks. On individual files, prefer to subset before chunking (suggestion 1). + +4. Chunk as early as possible, and avoid rechunking as much as possible. Always pass the ``chunks={}`` argument to :py:func:`~xarray.open_mfdataset` to avoid redundant file reads. -4. Using the h5netcdf package by passing ``engine='h5netcdf'`` to :py:meth:`~xarray.open_mfdataset` - can be quicker than the default ``engine='netcdf4'`` that uses the netCDF4 package. +5. Using the h5netcdf package by passing ``engine='h5netcdf'`` to :py:meth:`~xarray.open_mfdataset` can be quicker than the default ``engine='netcdf4'`` that uses the netCDF4 package. -5. Some dask-specific tips may be found `here `_. +6. Some dask-specific tips may be found `here `_. -6. The dask `diagnostics `_ can be - useful in identifying performance bottlenecks. +7. The dask `diagnostics `_ can be useful in identifying performance bottlenecks. diff --git a/doc/user-guide/interpolation.rst b/doc/user-guide/interpolation.rst index 73c0b312241..2dc47e9f591 100644 --- a/doc/user-guide/interpolation.rst +++ b/doc/user-guide/interpolation.rst @@ -132,8 +132,12 @@ It is now possible to safely compute the difference ``other - interpolated``. Interpolation methods --------------------- -We use :py:class:`scipy.interpolate.interp1d` for 1-dimensional interpolation and -:py:func:`scipy.interpolate.interpn` for multi-dimensional interpolation. +We use :py:class:`scipy.interpolate.interp1d` for 1-dimensional interpolation. +For multi-dimensional interpolation, an attempt is first made to decompose the +interpolation in a series of 1-dimensional interpolations, in which case +:py:class:`scipy.interpolate.interp1d` is used. If a decomposition cannot be +made (e.g. with advanced interpolation), :py:func:`scipy.interpolate.interpn` is +used. The interpolation method can be specified by the optional ``method`` argument. @@ -165,7 +169,9 @@ Additional keyword arguments can be passed to scipy's functions. [("time", np.arange(4)), ("space", [0.1, 0.2, 0.3])], ) - da.interp(time=4, space=np.linspace(-0.1, 0.5, 10), kwargs={"fill_value": None}) + da.interp( + time=4, space=np.linspace(-0.1, 0.5, 10), kwargs={"fill_value": "extrapolate"} + ) Advanced Interpolation @@ -198,23 +204,28 @@ For example: y = xr.DataArray([0.1, 0.2, 0.3], dims="z") da.sel(x=x, y=y) - # advanced interpolation - x = xr.DataArray([0.5, 1.5, 2.5], dims="z") - y = xr.DataArray([0.15, 0.25, 0.35], dims="z") + # advanced interpolation, without extrapolation + x = xr.DataArray([0.5, 1.5, 2.5, 3.5], dims="z") + y = xr.DataArray([0.15, 0.25, 0.35, 0.45], dims="z") da.interp(x=x, y=y) where values on the original coordinates -``(x, y) = ((0.5, 0.15), (1.5, 0.25), (2.5, 0.35))`` are obtained by the -2-dimensional interpolation and mapped along a new dimension ``z``. +``(x, y) = ((0.5, 0.15), (1.5, 0.25), (2.5, 0.35), (3.5, 0.45))`` are obtained +by the 2-dimensional interpolation and mapped along a new dimension ``z``. Since +no keyword arguments are passed to the interpolation routine, no extrapolation +is performed resulting in a ``nan`` value. If you want to add a coordinate to the new dimension ``z``, you can supply -:py:class:`~xarray.DataArray` s with a coordinate, +:py:class:`~xarray.DataArray` s with a coordinate. Extrapolation can be achieved +by passing additional arguments to SciPy's ``interpnd`` function, .. ipython:: python - x = xr.DataArray([0.5, 1.5, 2.5], dims="z", coords={"z": ["a", "b", "c"]}) - y = xr.DataArray([0.15, 0.25, 0.35], dims="z", coords={"z": ["a", "b", "c"]}) - da.interp(x=x, y=y) + x = xr.DataArray([0.5, 1.5, 2.5, 3.5], dims="z", coords={"z": ["a", "b", "c", "d"]}) + y = xr.DataArray( + [0.15, 0.25, 0.35, 0.45], dims="z", coords={"z": ["a", "b", "c", "d"]} + ) + da.interp(x=x, y=y, kwargs={"fill_value": None}) For the details of the advanced indexing, see :ref:`more advanced indexing `. diff --git a/doc/user-guide/io.rst b/doc/user-guide/io.rst index ddde0bf5888..81fa29bdf5f 100644 --- a/doc/user-guide/io.rst +++ b/doc/user-guide/io.rst @@ -518,8 +518,11 @@ the ability to store and analyze datasets far too large fit onto disk Xarray can't open just any zarr dataset, because xarray requires special metadata (attributes) describing the dataset dimensions and coordinates. -At this time, xarray can only open zarr datasets that have been written by -xarray. For implementation details, see :ref:`zarr_encoding`. +At this time, xarray can only open zarr datasets with these special attributes, +such as zarr datasets written by xarray, +`netCDF `_, +or `GDAL `_. +For implementation details, see :ref:`zarr_encoding`. To write a dataset with zarr, we use the :py:meth:`Dataset.to_zarr` method. @@ -548,6 +551,11 @@ store is already present at that path, an error will be raised, preventing it from being overwritten. To override this behavior and overwrite an existing store, add ``mode='w'`` when invoking :py:meth:`~Dataset.to_zarr`. +.. note:: + + xarray does not write NCZarr attributes. Therefore, NCZarr data must be + opened in read-only mode. + To store variable length strings, convert them to object arrays first with ``dtype=object``. diff --git a/doc/user-guide/plotting.rst b/doc/user-guide/plotting.rst index f514b4ecbef..78182ed265f 100644 --- a/doc/user-guide/plotting.rst +++ b/doc/user-guide/plotting.rst @@ -251,7 +251,7 @@ Finally, if a dataset does not have any coordinates it enumerates all data point .. ipython:: python :okwarning: - air1d_multi = air1d_multi.drop(["date", "time", "decimal_day"]) + air1d_multi = air1d_multi.drop_vars(["date", "time", "decimal_day"]) air1d_multi.plot() The same applies to 2D plots below. diff --git a/doc/user-guide/terminology.rst b/doc/user-guide/terminology.rst index 1876058323e..c8cfdd5133d 100644 --- a/doc/user-guide/terminology.rst +++ b/doc/user-guide/terminology.rst @@ -27,7 +27,7 @@ complete examples, please consult the relevant documentation.* Variable A `NetCDF-like variable - `_ + `_ consisting of dimensions, data, and attributes which describe a single array. The main functional difference between variables and numpy arrays is that numerical operations on variables implement array broadcasting diff --git a/doc/videos.yml b/doc/videos.yml new file mode 100644 index 00000000000..62c89563a56 --- /dev/null +++ b/doc/videos.yml @@ -0,0 +1,38 @@ +- title: "Xdev Python Tutorial Seminar Series 2022 Thinking with Xarray : High-level computation patterns" + src: '' + authors: + - Deepak Cherian +- title: "Xdev Python Tutorial Seminar Series 2021 seminar introducing xarray (2 of 2)" + src: '' + authors: + - Anderson Banihirwe + +- title: "Xdev Python Tutorial Seminar Series 2021 seminar introducing xarray (1 of 2)" + src: '' + authors: + - Anderson Banihirwe + +- title: "Xarray's 2020 virtual tutorial" + src: '' + authors: + - Anderson Banihirwe + - Deepak Cherian + - Martin Durant + +- title: "Xarray's Tutorial presented at the 2020 SciPy Conference" + src: ' ' + authors: + - Joe Hamman + - Deepak Cherian + - Ryan Abernathey + - Stephan Hoyer + +- title: "Scipy 2015 talk introducing xarray to a general audience" + src: '' + authors: + - Stephan Hoyer + +- title: " 2015 Unidata Users Workshop talk and tutorial with (`with answers`_) introducing xarray to users familiar with netCDF" + src: '' + authors: + - Stephan Hoyer diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 55de76bb9e7..54a273fbdc3 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,6 +22,8 @@ v2022.03.1 (unreleased) New Features ~~~~~~~~~~~~ +- The `zarr` backend is now able to read NCZarr. + By `Mattia Almansi `_. - Add a weighted ``quantile`` method to :py:class:`~core.weighted.DatasetWeighted` and :py:class:`~core.weighted.DataArrayWeighted` (:pull:`6059`). By `Christian Jauvin `_ and `David Huard `_. @@ -36,16 +38,56 @@ New Features elements which trigger summarization rather than full repr in (numpy) array detailed views of the html repr (:pull:`6400`). By `Benoît Bovy `_. +- Allow passing chunks in ``**kwargs`` form to :py:meth:`Dataset.chunk`, :py:meth:`DataArray.chunk`, and + :py:meth:`Variable.chunk`. (:pull:`6471`) + By `Tom Nicholas `_. +- Expose `inline_array` kwarg from `dask.array.from_array` in :py:func:`open_dataset`, :py:meth:`Dataset.chunk`, + :py:meth:`DataArray.chunk`, and :py:meth:`Variable.chunk`. (:pull:`6471`) + By `Tom Nicholas `_. +- :py:meth:`xr.polyval` now supports :py:class:`Dataset` and :py:class:`DataArray` args of any shape, + is faster and requires less memory. (:pull:`6548`) + By `Michael Niklas `_. +- Improved overall typing. +- :py:meth:`Dataset.to_dict` and :py:meth:`DataArray.to_dict` may now optionally include encoding + attributes. (:pull:`6635`) + By Joe Hamman `_. Breaking changes ~~~~~~~~~~~~~~~~ +- PyNIO support is now untested. The minimum versions of some dependencies were changed: + + =============== ===== ==== + Package Old New + =============== ===== ==== + cftime 1.2 1.4 + dask 2.30 2021.4 + distributed 2.30 2021.4 + h5netcdf 0.8 0.11 + matplotlib-base 3.3 3.4 + numba 0.51 0.53 + numpy 1.18 1.19 + pandas 1.1 1.2 + pint 0.16 0.17 + rasterio 1.1 1.2 + scipy 1.5 1.6 + sparse 0.11 0.12 + zarr 2.5 2.8 + =============== ===== ==== + - The Dataset and DataArray ``rename*`` methods do not implicitly add or drop indexes. (:pull:`5692`). By `Benoît Bovy `_. - Many arguments like ``keep_attrs``, ``axis``, and ``skipna`` are now keyword only for all reduction operations like ``.mean``. By `Deepak Cherian `_, `Jimmy Westling `_. +- Xarray's ufuncs have been removed, now that they can be replaced by numpy's ufuncs in all + supported versions of numpy. + By `Maximilian Roos `_. +- :py:meth:`xr.polyval` now uses the ``coord`` argument directly instead of its index coordinate. + (:pull:`6548`) + By `Michael Niklas `_. + Deprecations ~~~~~~~~~~~~ @@ -54,6 +96,8 @@ Deprecations Bug fixes ~~~~~~~~~ +- :py:meth:`Dataset.to_zarr` now allows to write all attribute types supported by `zarr-python`. + By `Mattia Almansi `_. - Set ``skipna=None`` for all ``quantile`` methods (e.g. :py:meth:`Dataset.quantile`) and ensure it skips missing values for float dtypes (consistent with other methods). This should not change the behavior (:pull:`6303`). @@ -62,22 +106,54 @@ Bug fixes coordinates. See the corresponding pull-request on GitHub for more details. (:pull:`5692`). By `Benoît Bovy `_. - Fixed "unhashable type" error trying to read NetCDF file with variable having its 'units' - attribute not ``str`` (e.g. ``numpy.ndarray``) (:issue:`6368`). - By `Oleh Khoma `_. + attribute not ``str`` (e.g. ``numpy.ndarray``) (:issue:`6368`). By `Oleh Khoma `_. +- Omit warning about specified dask chunks separating chunks on disk when the + underlying array is empty (e.g., because of an empty dimension) (:issue:`6401`). + By `Joseph K Aicher `_. - Fixed the poor html repr performance on large multi-indexes (:pull:`6400`). By `Benoît Bovy `_. - Allow fancy indexing of duck dask arrays along multiple dimensions. (:pull:`6414`) By `Justus Magin `_. +- In the API for backends, support dimensions that express their preferred chunk sizes + as a tuple of integers. (:issue:`6333`, :pull:`6334`) + By `Stan West `_. +- Fix bug in :py:func:`where` when passing non-xarray objects with ``keep_attrs=True``. (:issue:`6444`, :pull:`6461`) + By `Sam Levang `_. +- Allow passing both ``other`` and ``drop=True`` arguments to ``xr.DataArray.where`` + and ``xr.Dataset.where`` (:pull:`6466`, :pull:`6467`). + By `Michael Delgado `_. +- Ensure dtype encoding attributes are not added or modified on variables that + contain datetime-like values prior to being passed to + :py:func:`xarray.conventions.decode_cf_variable` (:issue:`6453`, + :pull:`6489`). By `Spencer Clark `_. +- Dark themes are now properly detected in Furo-themed Sphinx documents (:issue:`6500`, :pull:`6501`). + By `Kevin Paul `_. +- :py:meth:`isel` with `drop=True` works as intended with scalar :py:class:`DataArray` indexers. + (:issue:`6554`, :pull:`6579`) + By `Michael Niklas `_. +- Fixed silent overflow issue when decoding times encoded with 32-bit and below + unsigned integer data types (:issue:`6589`, :pull:`6598`). By `Spencer Clark + `_. Documentation ~~~~~~~~~~~~~ +- Revise the documentation for developers on specifying a backend's preferred chunk + sizes. In particular, correct the syntax and replace lists with tuples in the + examples. (:issue:`6333`, :pull:`6334`) + By `Stan West `_. + Performance ~~~~~~~~~~~ - GroupBy binary operations are now vectorized. Previously this involved looping over all groups. (:issue:`5804`,:pull:`6160`) By `Deepak Cherian `_. +- Substantially improved GroupBy operations using `flox `_. + This is auto-enabled when ``flox`` is installed. Use ``xr.set_options(use_flox=False)`` to use + the old algorithm. (:issue:`4473`, :issue:`4498`, :issue:`659`, :issue:`2237`, :pull:`271`). + By `Deepak Cherian `_,`Anderson Banihirwe `_, + `Jimmy Westling `_. Internal Changes ~~~~~~~~~~~~~~~~ diff --git a/setup.cfg b/setup.cfg index afa25325018..f5dd4dde810 100644 --- a/setup.cfg +++ b/setup.cfg @@ -75,8 +75,8 @@ zip_safe = False # https://mypy.readthedocs.io/en/latest/installed_packages.htm include_package_data = True python_requires = >=3.8 install_requires = - numpy >= 1.18 - pandas >= 1.1 + numpy >= 1.19 + pandas >= 1.2 packaging >= 20.0 [options.extras_require] @@ -98,7 +98,6 @@ accel = scipy bottleneck numbagg - numpy_groupies flox parallel = diff --git a/xarray/__init__.py b/xarray/__init__.py index aa9739d3d35..46dcf0e9b32 100644 --- a/xarray/__init__.py +++ b/xarray/__init__.py @@ -1,4 +1,4 @@ -from . import testing, tutorial, ufuncs +from . import testing, tutorial from .backends.api import ( load_dataarray, load_dataset, @@ -53,7 +53,6 @@ # `mypy --strict` running in projects that import xarray. __all__ = ( # Sub-packages - "ufuncs", "testing", "tutorial", # Top-level functions diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 548b98048ba..cc1e143e573 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1,23 +1,29 @@ +from __future__ import annotations + import os from glob import glob from io import BytesIO from numbers import Number from typing import ( TYPE_CHECKING, + Any, Callable, - Dict, + Final, Hashable, Iterable, + Literal, Mapping, MutableMapping, - Optional, - Tuple, + Sequence, + Type, Union, + cast, + overload, ) import numpy as np -from .. import backends, coding, conventions +from .. import backends, conventions from ..core import indexing from ..core.combine import ( _infer_concat_order_from_positions, @@ -26,6 +32,7 @@ ) from ..core.dataarray import DataArray from ..core.dataset import Dataset, _get_chunk, _maybe_chunk +from ..core.indexes import Index from ..core.utils import is_remote_uri from . import plugins from .common import AbstractDataStore, ArrayWriter, _normalize_path @@ -35,7 +42,27 @@ try: from dask.delayed import Delayed except ImportError: - Delayed = None + Delayed = None # type: ignore + from ..core.types import ( + CombineAttrsOptions, + CompatOptions, + JoinOptions, + NestedSequence, + ) + from .common import BackendEntrypoint + + T_NetcdfEngine = Literal["netcdf4", "scipy", "h5netcdf"] + T_Engine = Union[ + T_NetcdfEngine, + Literal["pydap", "pynio", "pseudonetcdf", "cfgrib", "zarr"], + Type[BackendEntrypoint], + str, # no nice typing support for custom backends + None, + ] + T_Chunks = Union[int, dict[Any, Any], Literal["auto"], None] + T_NetcdfTypes = Literal[ + "NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_64BIT", "NETCDF3_CLASSIC" + ] DATAARRAY_NAME = "__xarray_dataarray_name__" @@ -53,7 +80,8 @@ } -def _get_default_engine_remote_uri(): +def _get_default_engine_remote_uri() -> Literal["netcdf4", "pydap"]: + engine: Literal["netcdf4", "pydap"] try: import netCDF4 # noqa: F401 @@ -71,17 +99,18 @@ def _get_default_engine_remote_uri(): return engine -def _get_default_engine_gz(): +def _get_default_engine_gz() -> Literal["scipy"]: try: import scipy # noqa: F401 - engine = "scipy" + engine: Final = "scipy" except ImportError: # pragma: no cover raise ValueError("scipy is required for accessing .gz files") return engine -def _get_default_engine_netcdf(): +def _get_default_engine_netcdf() -> Literal["netcdf4", "scipy"]: + engine: Literal["netcdf4", "scipy"] try: import netCDF4 # noqa: F401 @@ -99,19 +128,19 @@ def _get_default_engine_netcdf(): return engine -def _get_default_engine(path: str, allow_remote: bool = False): +def _get_default_engine(path: str, allow_remote: bool = False) -> T_NetcdfEngine: if allow_remote and is_remote_uri(path): - return _get_default_engine_remote_uri() + return _get_default_engine_remote_uri() # type: ignore[return-value] elif path.endswith(".gz"): return _get_default_engine_gz() else: return _get_default_engine_netcdf() -def _validate_dataset_names(dataset): +def _validate_dataset_names(dataset: Dataset) -> None: """DataArray.name and Dataset keys must be a string or None""" - def check_name(name): + def check_name(name: Hashable): if isinstance(name, str): if not name: raise ValueError( @@ -216,7 +245,7 @@ def _finalize_store(write, store): store.close() -def load_dataset(filename_or_obj, **kwargs): +def load_dataset(filename_or_obj, **kwargs) -> Dataset: """Open, load into memory, and close a Dataset from a file or file-like object. @@ -274,6 +303,7 @@ def _chunk_ds( engine, chunks, overwrite_encoded_chunks, + inline_array, **extra_tokens, ): from dask.base import tokenize @@ -292,6 +322,7 @@ def _chunk_ds( overwrite_encoded_chunks=overwrite_encoded_chunks, name_prefix=name_prefix, token=token, + inline_array=inline_array, ) return backend_ds._replace(variables) @@ -303,6 +334,7 @@ def _dataset_from_backend_dataset( chunks, cache, overwrite_encoded_chunks, + inline_array, **extra_tokens, ): if not isinstance(chunks, (int, dict)) and chunks not in {None, "auto"}: @@ -320,6 +352,7 @@ def _dataset_from_backend_dataset( engine, chunks, overwrite_encoded_chunks, + inline_array, **extra_tokens, ) @@ -333,22 +366,23 @@ def _dataset_from_backend_dataset( def open_dataset( - filename_or_obj, - *args, - engine=None, - chunks=None, - cache=None, - decode_cf=None, - mask_and_scale=None, - decode_times=None, - decode_timedelta=None, - use_cftime=None, - concat_characters=None, - decode_coords=None, - drop_variables=None, - backend_kwargs=None, + filename_or_obj: str | os.PathLike, + *, + engine: T_Engine = None, + chunks: T_Chunks = None, + cache: bool | None = None, + decode_cf: bool | None = None, + mask_and_scale: bool | None = None, + decode_times: bool | None = None, + decode_timedelta: bool | None = None, + use_cftime: bool | None = None, + concat_characters: bool | None = None, + decode_coords: Literal["coordinates", "all"] | bool | None = None, + drop_variables: str | Iterable[str] | None = None, + inline_array: bool = False, + backend_kwargs: dict[str, Any] | None = None, **kwargs, -): +) -> Dataset: """Open and decode a dataset from a file or file-like object. Parameters @@ -360,12 +394,13 @@ def open_dataset( scipy.io.netcdf (only netCDF3 supported). Byte-strings or file-like objects are opened by scipy.io.netcdf (netCDF3) or h5py (netCDF4/HDF). engine : {"netcdf4", "scipy", "pydap", "h5netcdf", "pynio", "cfgrib", \ - "pseudonetcdf", "zarr"} or subclass of xarray.backends.BackendEntrypoint, optional + "pseudonetcdf", "zarr", None}, installed backend \ + or subclass of xarray.backends.BackendEntrypoint, optional Engine to use when reading files. If not provided, the default engine is chosen based on available dependencies, with a preference for "netcdf4". A custom backend class (a subclass of ``BackendEntrypoint``) can also be used. - chunks : int or dict, optional + chunks : int, dict, 'auto' or None, optional If chunks is provided, it is used to load the new dataset into dask arrays. ``chunks=-1`` loads the dataset with dask using a single chunk for all arrays. ``chunks={}`` loads the dataset with dask using @@ -426,10 +461,16 @@ def open_dataset( as coordinate variables. - "all": Set variables referred to in ``'grid_mapping'``, ``'bounds'`` and other attributes as coordinate variables. - drop_variables: str or iterable, optional + drop_variables: str or iterable of str, optional A variable or list of variables to exclude from being parsed from the dataset. This may be useful to drop variables with problems or inconsistent values. + inline_array: bool, default: False + How to include the array in the dask task graph. + By default(``inline_array=False``) the array is included in a task by + itself, and each chunk refers to that task by its key. With + ``inline_array=True``, Dask will instead inline the array directly + in the values of the task graph. See :py:func:`dask.array.from_array`. backend_kwargs: dict Additional keyword arguments passed on to the engine open function, equivalent to `**kwargs`. @@ -463,11 +504,6 @@ def open_dataset( -------- open_mfdataset """ - if len(args) > 0: - raise TypeError( - "open_dataset() takes only 1 positional argument starting from version 0.18.0, " - "all other options must be passed as keyword arguments" - ) if cache is None: cache = chunks is None @@ -505,6 +541,7 @@ def open_dataset( chunks, cache, overwrite_encoded_chunks, + inline_array, drop_variables=drop_variables, **decoders, **kwargs, @@ -513,22 +550,23 @@ def open_dataset( def open_dataarray( - filename_or_obj, - *args, - engine=None, - chunks=None, - cache=None, - decode_cf=None, - mask_and_scale=None, - decode_times=None, - decode_timedelta=None, - use_cftime=None, - concat_characters=None, - decode_coords=None, - drop_variables=None, - backend_kwargs=None, + filename_or_obj: str | os.PathLike, + *, + engine: T_Engine = None, + chunks: T_Chunks = None, + cache: bool | None = None, + decode_cf: bool | None = None, + mask_and_scale: bool | None = None, + decode_times: bool | None = None, + decode_timedelta: bool | None = None, + use_cftime: bool | None = None, + concat_characters: bool | None = None, + decode_coords: Literal["coordinates", "all"] | bool | None = None, + drop_variables: str | Iterable[str] | None = None, + inline_array: bool = False, + backend_kwargs: dict[str, Any] | None = None, **kwargs, -): +) -> DataArray: """Open an DataArray from a file or file-like object containing a single data variable. @@ -544,11 +582,12 @@ def open_dataarray( scipy.io.netcdf (only netCDF3 supported). Byte-strings or file-like objects are opened by scipy.io.netcdf (netCDF3) or h5py (netCDF4/HDF). engine : {"netcdf4", "scipy", "pydap", "h5netcdf", "pynio", "cfgrib", \ - "pseudonetcdf", "zarr"}, optional + "pseudonetcdf", "zarr", None}, installed backend \ + or subclass of xarray.backends.BackendEntrypoint, optional Engine to use when reading files. If not provided, the default engine is chosen based on available dependencies, with a preference for "netcdf4". - chunks : int or dict, optional + chunks : int, dict, 'auto' or None, optional If chunks is provided, it is used to load the new dataset into dask arrays. ``chunks=-1`` loads the dataset with dask using a single chunk for all arrays. `chunks={}`` loads the dataset with dask using @@ -609,10 +648,16 @@ def open_dataarray( as coordinate variables. - "all": Set variables referred to in ``'grid_mapping'``, ``'bounds'`` and other attributes as coordinate variables. - drop_variables: str or iterable, optional + drop_variables: str or iterable of str, optional A variable or list of variables to exclude from being parsed from the dataset. This may be useful to drop variables with problems or inconsistent values. + inline_array: bool, default: False + How to include the array in the dask task graph. + By default(``inline_array=False``) the array is included in a task by + itself, and each chunk refers to that task by its key. With + ``inline_array=True``, Dask will instead inline the array directly + in the values of the task graph. See :py:func:`dask.array.from_array`. backend_kwargs: dict Additional keyword arguments passed on to the engine open function, equivalent to `**kwargs`. @@ -643,11 +688,6 @@ def open_dataarray( -------- open_dataset """ - if len(args) > 0: - raise TypeError( - "open_dataarray() takes only 1 positional argument starting from version 0.18.0, " - "all other options must be passed as keyword arguments" - ) dataset = open_dataset( filename_or_obj, @@ -660,6 +700,7 @@ def open_dataarray( chunks=chunks, cache=cache, drop_variables=drop_variables, + inline_array=inline_array, backend_kwargs=backend_kwargs, use_cftime=use_cftime, decode_timedelta=decode_timedelta, @@ -690,21 +731,27 @@ def open_dataarray( def open_mfdataset( - paths, - chunks=None, - concat_dim=None, - compat="no_conflicts", - preprocess=None, - engine=None, - data_vars="all", + paths: str | NestedSequence[str | os.PathLike], + chunks: T_Chunks = None, + concat_dim: str + | DataArray + | Index + | Sequence[str] + | Sequence[DataArray] + | Sequence[Index] + | None = None, + compat: CompatOptions = "no_conflicts", + preprocess: Callable[[Dataset], Dataset] | None = None, + engine: T_Engine = None, + data_vars: Literal["all", "minimal", "different"] | list[str] = "all", coords="different", - combine="by_coords", - parallel=False, - join="outer", - attrs_file=None, - combine_attrs="override", + combine: Literal["by_coords", "nested"] = "by_coords", + parallel: bool = False, + join: JoinOptions = "outer", + attrs_file: str | os.PathLike | None = None, + combine_attrs: CombineAttrsOptions = "override", **kwargs, -): +) -> Dataset: """Open multiple files as a single dataset. If combine='by_coords' then the function ``combine_by_coords`` is used to combine @@ -718,19 +765,19 @@ def open_mfdataset( Parameters ---------- - paths : str or sequence + paths : str or nested sequence of paths Either a string glob in the form ``"path/to/my/files/*.nc"`` or an explicit list of files to open. Paths can be given as strings or as pathlib Paths. If concatenation along more than one dimension is desired, then ``paths`` must be a nested list-of-lists (see ``combine_nested`` for details). (A string glob will be expanded to a 1-dimensional list.) - chunks : int or dict, optional + chunks : int, dict, 'auto' or None, optional Dictionary with keys given by dimension names and values given by chunk sizes. In general, these should divide the dimensions of each dataset. If int, chunk each dimension by ``chunks``. By default, chunks will be chosen to load entire input files into memory at once. This has a major impact on performance: please see the full documentation for more details [2]_. - concat_dim : str, or list of str, DataArray, Index or None, optional + concat_dim : str, DataArray, Index or a Sequence of these or None, optional Dimensions to concatenate files along. You only need to provide this argument if ``combine='nested'``, and if any of the dimensions along which you want to concatenate is not a dimension in the original datasets, e.g., if you want to @@ -743,7 +790,7 @@ def open_mfdataset( Whether ``xarray.combine_by_coords`` or ``xarray.combine_nested`` is used to combine all the data. Default is to use ``xarray.combine_by_coords``. compat : {"identical", "equals", "broadcast_equals", \ - "no_conflicts", "override"}, optional + "no_conflicts", "override"}, default: "no_conflicts" String indicating how to compare variables of the same name for potential conflicts when merging: @@ -761,12 +808,13 @@ def open_mfdataset( If provided, call this function on each dataset prior to concatenation. You can find the file-name from which each dataset was loaded in ``ds.encoding["source"]``. - engine : {"netcdf4", "scipy", "pydap", "h5netcdf", "pynio", "cfgrib", "zarr"}, \ - optional + engine : {"netcdf4", "scipy", "pydap", "h5netcdf", "pynio", "cfgrib", \ + "pseudonetcdf", "zarr", None}, installed backend \ + or subclass of xarray.backends.BackendEntrypoint, optional Engine to use when reading files. If not provided, the default engine is chosen based on available dependencies, with a preference for "netcdf4". - data_vars : {"minimal", "different", "all"} or list of str, optional + data_vars : {"minimal", "different", "all"} or list of str, default: "all" These data variables will be concatenated together: * "minimal": Only data variables in which the dimension already appears are included. @@ -791,10 +839,10 @@ def open_mfdataset( those corresponding to other dimensions. * list of str: The listed coordinate variables will be concatenated, in addition the "minimal" coordinates. - parallel : bool, optional + parallel : bool, default: False If True, the open and preprocess steps of this function will be performed in parallel using ``dask.delayed``. Default is False. - join : {"outer", "inner", "left", "right", "exact, "override"}, optional + join : {"outer", "inner", "left", "right", "exact", "override"}, default: "outer" String indicating how to combine differing indexes (excluding concat_dim) in objects @@ -811,6 +859,22 @@ def open_mfdataset( Path of the file used to read global attributes from. By default global attributes are read from the first file provided, with wildcard matches sorted by filename. + combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \ + "override"} or callable, default: "override" + A callable or a string indicating how to combine attrs of the objects being + merged: + + - "drop": empty attrs on returned Dataset. + - "identical": all attrs must be the same on every object. + - "no_conflicts": attrs from all objects are combined, any that have + the same name must also have the same value. + - "drop_conflicts": attrs from all objects are combined, any that have + the same name but different values are dropped. + - "override": skip comparing and copy attrs from the first dataset to + the result. + + If a callable, it must expect a sequence of ``attrs`` dicts and a context object + as its only parameters. **kwargs : optional Additional arguments passed on to :py:func:`xarray.open_dataset`. @@ -854,8 +918,8 @@ def open_mfdataset( ), expand=False, ) - paths = fs.glob(fs._strip_protocol(paths)) # finds directories - paths = [fs.get_mapper(path) for path in paths] + tmp_paths = fs.glob(fs._strip_protocol(paths)) # finds directories + paths = [fs.get_mapper(path) for path in tmp_paths] elif is_remote_uri(paths): raise ValueError( "cannot do wild-card matching for paths that are remote URLs " @@ -874,7 +938,7 @@ def open_mfdataset( if combine == "nested": if isinstance(concat_dim, (str, DataArray)) or concat_dim is None: - concat_dim = [concat_dim] + concat_dim = [concat_dim] # type: ignore[assignment] # This creates a flat list which is easier to iterate over, whilst # encoding the originally-supplied structure as "ids". @@ -960,32 +1024,106 @@ def multi_file_closer(): # read global attributes from the attrs_file or from the first dataset if attrs_file is not None: if isinstance(attrs_file, os.PathLike): - attrs_file = os.fspath(attrs_file) + attrs_file = cast(str, os.fspath(attrs_file)) combined.attrs = datasets[paths.index(attrs_file)].attrs return combined -WRITEABLE_STORES: Dict[str, Callable] = { +WRITEABLE_STORES: dict[T_NetcdfEngine, Callable] = { "netcdf4": backends.NetCDF4DataStore.open, "scipy": backends.ScipyDataStore, "h5netcdf": backends.H5NetCDFStore.open, } +# multifile=True returns writer and datastore +@overload +def to_netcdf( + dataset: Dataset, + path_or_file: str | os.PathLike | None = None, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: bool = True, + *, + multifile: Literal[True], + invalid_netcdf: bool = False, +) -> tuple[ArrayWriter, AbstractDataStore]: + ... + + +# path=None writes to bytes +@overload +def to_netcdf( + dataset: Dataset, + path_or_file: None = None, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: bool = True, + multifile: Literal[False] = False, + invalid_netcdf: bool = False, +) -> bytes: + ... + + +# compute=False returns dask.Delayed +@overload +def to_netcdf( + dataset: Dataset, + path_or_file: str | os.PathLike, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + *, + compute: Literal[False], + multifile: Literal[False] = False, + invalid_netcdf: bool = False, +) -> Delayed: + ... + + +# default return None +@overload def to_netcdf( dataset: Dataset, - path_or_file=None, - mode: str = "w", - format: str = None, - group: str = None, - engine: str = None, - encoding: Mapping = None, - unlimited_dims: Iterable[Hashable] = None, + path_or_file: str | os.PathLike, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: Literal[True] = True, + multifile: Literal[False] = False, + invalid_netcdf: bool = False, +) -> None: + ... + + +def to_netcdf( + dataset: Dataset, + path_or_file: str | os.PathLike | None = None, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, compute: bool = True, multifile: bool = False, invalid_netcdf: bool = False, -) -> Union[Tuple[ArrayWriter, AbstractDataStore], bytes, "Delayed", None]: +) -> tuple[ArrayWriter, AbstractDataStore] | bytes | Delayed | None: """This function creates an appropriate datastore for writing a dataset to disk as a netCDF file @@ -1030,7 +1168,7 @@ def to_netcdf( raise ValueError(f"unrecognized engine for to_netcdf: {engine!r}") if format is not None: - format = format.upper() + format = format.upper() # type: ignore[assignment] # handle scheduler specific logic scheduler = _get_scheduler() @@ -1080,7 +1218,7 @@ def to_netcdf( writes = writer.sync(compute=compute) - if path_or_file is None: + if isinstance(target, BytesIO): store.sync() return target.getvalue() finally: @@ -1273,49 +1411,102 @@ def _validate_region(ds, region): f"{list(region.keys())}, but that is not " f"the case for some variables here. To drop these variables " f"from this dataset before exporting to zarr, write: " - f".drop({non_matching_vars!r})" + f".drop_vars({non_matching_vars!r})" ) -def _validate_datatypes_for_zarr_append(dataset): - """DataArray.name and Dataset keys must be a string or None""" +def _validate_datatypes_for_zarr_append(zstore, dataset): + """If variable exists in the store, confirm dtype of the data to append is compatible with + existing dtype. + """ + + existing_vars = zstore.get_variables() - def check_dtype(var): + def check_dtype(vname, var): if ( - not np.issubdtype(var.dtype, np.number) - and not np.issubdtype(var.dtype, np.datetime64) - and not np.issubdtype(var.dtype, np.bool_) - and not coding.strings.is_unicode_dtype(var.dtype) - and not var.dtype == object + vname not in existing_vars + or np.issubdtype(var.dtype, np.number) + or np.issubdtype(var.dtype, np.datetime64) + or np.issubdtype(var.dtype, np.bool_) + or var.dtype == object ): - # and not re.match('^bytes[1-9]+$', var.dtype.name)): + # We can skip dtype equality checks under two conditions: (1) if the var to append is + # new to the dataset, because in this case there is no existing var to compare it to; + # or (2) if var to append's dtype is known to be easy-to-append, because in this case + # we can be confident appending won't cause problems. Examples of dtypes which are not + # easy-to-append include length-specified strings of type `|S*` or ` backends.ZarrStore: + ... + +# compute=False returns dask.Delayed +@overload def to_zarr( dataset: Dataset, - store: Union[MutableMapping, str, os.PathLike] = None, - chunk_store=None, - mode: str = None, + store: MutableMapping | str | os.PathLike[str] | None = None, + chunk_store: MutableMapping | str | os.PathLike | None = None, + mode: Literal["w", "w-", "a", "r+", None] = None, synchronizer=None, - group: str = None, - encoding: Mapping = None, + group: str | None = None, + encoding: Mapping | None = None, + *, + compute: Literal[False], + consolidated: bool | None = None, + append_dim: Hashable | None = None, + region: Mapping[str, slice] | None = None, + safe_chunks: bool = True, + storage_options: dict[str, str] | None = None, +) -> Delayed: + ... + + +def to_zarr( + dataset: Dataset, + store: MutableMapping | str | os.PathLike[str] | None = None, + chunk_store: MutableMapping | str | os.PathLike | None = None, + mode: Literal["w", "w-", "a", "r+", None] = None, + synchronizer=None, + group: str | None = None, + encoding: Mapping | None = None, compute: bool = True, - consolidated: Optional[bool] = None, - append_dim: Hashable = None, - region: Mapping[str, slice] = None, + consolidated: bool | None = None, + append_dim: Hashable | None = None, + region: Mapping[str, slice] | None = None, safe_chunks: bool = True, - storage_options: Dict[str, str] = None, -): + storage_options: dict[str, str] | None = None, +) -> backends.ZarrStore | Delayed: """This function creates an appropriate datastore for writing a dataset to a zarr ztore @@ -1370,9 +1561,8 @@ def to_zarr( f"'w-', 'a' and 'r+', but mode={mode!r}" ) - # validate Dataset keys, DataArray names, and attr keys/values + # validate Dataset keys, DataArray names _validate_dataset_names(dataset) - _validate_attrs(dataset) if region is not None: _validate_region(dataset, region) @@ -1403,7 +1593,7 @@ def to_zarr( ) if mode in ["a", "r+"]: - _validate_datatypes_for_zarr_append(dataset) + _validate_datatypes_for_zarr_append(zstore, dataset) if append_dim is not None: existing_dims = zstore.get_dimensions() if append_dim not in existing_dims: diff --git a/xarray/backends/common.py b/xarray/backends/common.py index ad92a6c5869..52738c639e1 100644 --- a/xarray/backends/common.py +++ b/xarray/backends/common.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import logging import os import time import traceback -from typing import Any, Dict, Tuple, Type, Union +from typing import Any import numpy as np @@ -369,13 +371,13 @@ class BackendEntrypoint: method is not mandatory. """ - open_dataset_parameters: Union[Tuple, None] = None + open_dataset_parameters: tuple | None = None """list of ``open_dataset`` method parameters""" def open_dataset( self, - filename_or_obj: str, - drop_variables: Tuple[str] = None, + filename_or_obj: str | os.PathLike, + drop_variables: tuple[str] | None = None, **kwargs: Any, ): """ @@ -384,7 +386,7 @@ def open_dataset( raise NotImplementedError - def guess_can_open(self, filename_or_obj): + def guess_can_open(self, filename_or_obj: str | os.PathLike): """ Backend open_dataset method used by Xarray in :py:func:`~xarray.open_dataset`. """ @@ -392,4 +394,4 @@ def guess_can_open(self, filename_or_obj): return False -BACKEND_ENTRYPOINTS: Dict[str, Type[BackendEntrypoint]] = {} +BACKEND_ENTRYPOINTS: dict[str, type[BackendEntrypoint]] = {} diff --git a/xarray/backends/locks.py b/xarray/backends/locks.py index 59417336f5f..1cc93779843 100644 --- a/xarray/backends/locks.py +++ b/xarray/backends/locks.py @@ -7,12 +7,12 @@ from dask.utils import SerializableLock except ImportError: # no need to worry about serializing the lock - SerializableLock = threading.Lock + SerializableLock = threading.Lock # type: ignore try: from dask.distributed import Lock as DistributedLock except ImportError: - DistributedLock = None + DistributedLock = None # type: ignore # Locks used by multiple backends. diff --git a/xarray/backends/plugins.py b/xarray/backends/plugins.py index 7444fbf11eb..44953b875d9 100644 --- a/xarray/backends/plugins.py +++ b/xarray/backends/plugins.py @@ -1,6 +1,7 @@ import functools import inspect import itertools +import sys import warnings from importlib.metadata import entry_points @@ -95,7 +96,11 @@ def build_engines(entrypoints): @functools.lru_cache(maxsize=1) def list_engines(): - entrypoints = entry_points().get("xarray.backends", ()) + # New selection mechanism introduced with Python 3.10. See GH6514. + if sys.version_info >= (3, 10): + entrypoints = entry_points(group="xarray.backends") + else: + entrypoints = entry_points().get("xarray.backends", ()) return build_engines(entrypoints) diff --git a/xarray/backends/scipy_.py b/xarray/backends/scipy_.py index 4c1ce1ef09d..df3ee364546 100644 --- a/xarray/backends/scipy_.py +++ b/xarray/backends/scipy_.py @@ -86,7 +86,8 @@ def _open_scipy_netcdf(filename, mode, mmap, version): ) except TypeError as e: # TODO: gzipped loading only works with NetCDF3 files. - if "is not a valid NetCDF 3 file" in e.message: + errmsg = e.args[0] + if "is not a valid NetCDF 3 file" in errmsg: raise ValueError("gzipped file loading only supports NetCDF 3 files.") else: raise diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index aca0b8064f5..1b8b7ee81e7 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -1,3 +1,4 @@ +import json import os import warnings @@ -178,19 +179,37 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name, safe_chunks): raise AssertionError("We should never get here. Function logic must be wrong.") -def _get_zarr_dims_and_attrs(zarr_obj, dimension_key): +def _get_zarr_dims_and_attrs(zarr_obj, dimension_key, try_nczarr): # Zarr arrays do not have dimensions. To get around this problem, we add # an attribute that specifies the dimension. We have to hide this attribute # when we send the attributes to the user. # zarr_obj can be either a zarr group or zarr array try: + # Xarray-Zarr dimensions = zarr_obj.attrs[dimension_key] - except KeyError: - raise KeyError( - f"Zarr object is missing the attribute `{dimension_key}`, which is " - "required for xarray to determine variable dimensions." - ) - attributes = HiddenKeyDict(zarr_obj.attrs, [dimension_key]) + except KeyError as e: + if not try_nczarr: + raise KeyError( + f"Zarr object is missing the attribute `{dimension_key}`, which is " + "required for xarray to determine variable dimensions." + ) from e + + # NCZarr defines dimensions through metadata in .zarray + zarray_path = os.path.join(zarr_obj.path, ".zarray") + zarray = json.loads(zarr_obj.store[zarray_path]) + try: + # NCZarr uses Fully Qualified Names + dimensions = [ + os.path.basename(dim) for dim in zarray["_NCZARR_ARRAY"]["dimrefs"] + ] + except KeyError as e: + raise KeyError( + f"Zarr object is missing the attribute `{dimension_key}` and the NCZarr metadata, " + "which are required for xarray to determine variable dimensions." + ) from e + + nc_attrs = [attr for attr in zarr_obj.attrs if attr.startswith("_NC")] + attributes = HiddenKeyDict(zarr_obj.attrs, [dimension_key] + nc_attrs) return dimensions, attributes @@ -301,6 +320,15 @@ def _validate_existing_dims(var_name, new_var, existing_var, region, append_dim) ) +def _put_attrs(zarr_obj, attrs): + """Raise a more informative error message for invalid attrs.""" + try: + zarr_obj.attrs.put(attrs) + except TypeError as e: + raise TypeError("Invalid attribute in Dataset.attrs.") from e + return zarr_obj + + class ZarrStore(AbstractWritableDataStore): """Store for reading and writing data via zarr""" @@ -409,7 +437,10 @@ def ds(self): def open_store_variable(self, name, zarr_array): data = indexing.LazilyIndexedArray(ZarrArrayWrapper(name, self)) - dimensions, attributes = _get_zarr_dims_and_attrs(zarr_array, DIMENSION_KEY) + try_nczarr = self._mode == "r" + dimensions, attributes = _get_zarr_dims_and_attrs( + zarr_array, DIMENSION_KEY, try_nczarr + ) attributes = dict(attributes) encoding = { "chunks": zarr_array.chunks, @@ -430,26 +461,24 @@ def get_variables(self): ) def get_attrs(self): - return dict(self.zarr_group.attrs.asdict()) + return { + k: v + for k, v in self.zarr_group.attrs.asdict().items() + if not k.startswith("_NC") + } def get_dimensions(self): + try_nczarr = self._mode == "r" dimensions = {} for k, v in self.zarr_group.arrays(): - try: - for d, s in zip(v.attrs[DIMENSION_KEY], v.shape): - if d in dimensions and dimensions[d] != s: - raise ValueError( - f"found conflicting lengths for dimension {d} " - f"({s} != {dimensions[d]})" - ) - dimensions[d] = s - - except KeyError: - raise KeyError( - f"Zarr object is missing the attribute `{DIMENSION_KEY}`, " - "which is required for xarray to determine " - "variable dimensions." - ) + dim_names, _ = _get_zarr_dims_and_attrs(v, DIMENSION_KEY, try_nczarr) + for d, s in zip(dim_names, v.shape): + if d in dimensions and dimensions[d] != s: + raise ValueError( + f"found conflicting lengths for dimension {d} " + f"({s} != {dimensions[d]})" + ) + dimensions[d] = s return dimensions def set_dimensions(self, variables, unlimited_dims=None): @@ -459,7 +488,7 @@ def set_dimensions(self, variables, unlimited_dims=None): ) def set_attributes(self, attributes): - self.zarr_group.attrs.put(attributes) + _put_attrs(self.zarr_group, attributes) def encode_variable(self, variable): variable = encode_zarr_variable(variable) @@ -598,7 +627,7 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No zarr_array = self.zarr_group.create( name, shape=shape, dtype=dtype, fill_value=fill_value, **encoding ) - zarr_array.attrs.put(encoded_attrs) + zarr_array = _put_attrs(zarr_array, encoded_attrs) write_region = self._write_region if self._write_region is not None else {} write_region = {dim: write_region.get(dim, slice(None)) for dim in dims} @@ -645,7 +674,7 @@ def open_zarr( The `store` object should be a valid store for a Zarr group. `store` variables must contain dimension metadata encoded in the - `_ARRAY_DIMENSIONS` attribute. + `_ARRAY_DIMENSIONS` attribute or must have NCZarr format. Parameters ---------- diff --git a/xarray/coding/times.py b/xarray/coding/times.py index 42a815300e5..eb8c1dbc42a 100644 --- a/xarray/coding/times.py +++ b/xarray/coding/times.py @@ -2,6 +2,7 @@ import warnings from datetime import datetime, timedelta from functools import partial +from typing import TYPE_CHECKING import numpy as np import pandas as pd @@ -27,6 +28,9 @@ except ImportError: cftime = None +if TYPE_CHECKING: + from ..core.types import CFCalendar + # standard calendars recognized by cftime _STANDARD_CALENDARS = {"standard", "gregorian", "proleptic_gregorian"} @@ -218,9 +222,12 @@ def _decode_datetime_with_pandas(flat_num_dates, units, calendar): pd.to_timedelta(flat_num_dates.max(), delta) + ref_date # To avoid integer overflow when converting to nanosecond units for integer - # dtypes smaller than np.int64 cast all integer-dtype arrays to np.int64 - # (GH 2002). - if flat_num_dates.dtype.kind == "i": + # dtypes smaller than np.int64 cast all integer and unsigned integer dtype + # arrays to np.int64 (GH 2002, GH 6589). Note this is safe even in the case + # of np.uint64 values, because any np.uint64 value that would lead to + # overflow when converting to np.int64 would not be representable with a + # timedelta64 value, and therefore would raise an error in the lines above. + if flat_num_dates.dtype.kind in "iu": flat_num_dates = flat_num_dates.astype(np.int64) # Cast input ordinals to integers of nanoseconds because pd.to_timedelta @@ -341,7 +348,7 @@ def _infer_time_units_from_diff(unique_timedeltas): return "seconds" -def infer_calendar_name(dates): +def infer_calendar_name(dates) -> "CFCalendar": """Given an array of datetimes, infer the CF calendar name""" if is_np_datetime_like(dates.dtype): return "proleptic_gregorian" diff --git a/xarray/conventions.py b/xarray/conventions.py index ae915069947..102ef003186 100644 --- a/xarray/conventions.py +++ b/xarray/conventions.py @@ -7,7 +7,7 @@ from .coding import strings, times, variables from .coding.variables import SerializationWarning, pop_to from .core import duck_array_ops, indexing -from .core.common import contains_cftime_datetimes +from .core.common import _contains_datetime_like_objects, contains_cftime_datetimes from .core.pycompat import is_duck_dask_array from .core.variable import IndexVariable, Variable, as_variable @@ -340,6 +340,11 @@ def decode_cf_variable( A variable holding the decoded equivalent of var. """ var = as_variable(var) + + # Ensure datetime-like Variables are passed through unmodified (GH 6453) + if _contains_datetime_like_objects(var): + return var + original_dtype = var.dtype if decode_timedelta is None: @@ -770,7 +775,7 @@ def _encode_coordinates(variables, attributes, non_dim_coord_names): # this will copy coordinates from encoding to attrs if "coordinates" in attrs # after the next line, "coordinates" is never in encoding # we get support for attrs["coordinates"] for free. - coords_str = pop_to(encoding, attrs, "coordinates") + coords_str = pop_to(encoding, attrs, "coordinates") or attrs.get("coordinates") if not coords_str and variable_coordinates[name]: coordinates_text = " ".join( str(coord_name) diff --git a/xarray/core/_reductions.py b/xarray/core/_reductions.py index 7df2fc16746..d782363760a 100644 --- a/xarray/core/_reductions.py +++ b/xarray/core/_reductions.py @@ -14,7 +14,7 @@ try: import flox except ImportError: - flox = None + flox = None # type: ignore class DatasetReductions: diff --git a/xarray/core/_typed_ops.pyi b/xarray/core/_typed_ops.pyi index e23b5848ff7..e5b3c9112c7 100644 --- a/xarray/core/_typed_ops.pyi +++ b/xarray/core/_typed_ops.pyi @@ -21,7 +21,7 @@ from .variable import Variable try: from dask.array import Array as DaskArray except ImportError: - DaskArray = np.ndarray + DaskArray = np.ndarray # type: ignore # DatasetOpsMixin etc. are parent classes of Dataset etc. # Because of https://github.com/pydata/xarray/issues/5755, we redefine these. Generally diff --git a/xarray/core/accessor_dt.py b/xarray/core/accessor_dt.py index 7f8bf79a50a..c90ad204a4a 100644 --- a/xarray/core/accessor_dt.py +++ b/xarray/core/accessor_dt.py @@ -1,4 +1,7 @@ +from __future__ import annotations + import warnings +from typing import TYPE_CHECKING, Generic import numpy as np import pandas as pd @@ -11,6 +14,12 @@ ) from .npcompat import DTypeLike from .pycompat import is_duck_dask_array +from .types import T_DataArray + +if TYPE_CHECKING: + from .dataarray import DataArray + from .dataset import Dataset + from .types import CFCalendar def _season_from_months(months): @@ -156,7 +165,7 @@ def _round_field(values, name, freq): return _round_through_series_or_index(values, name, freq) -def _strftime_through_cftimeindex(values, date_format): +def _strftime_through_cftimeindex(values, date_format: str): """Coerce an array of cftime-like values to a CFTimeIndex and access requested datetime component """ @@ -168,7 +177,7 @@ def _strftime_through_cftimeindex(values, date_format): return field_values.values.reshape(values.shape) -def _strftime_through_series(values, date_format): +def _strftime_through_series(values, date_format: str): """Coerce an array of datetime-like values to a pandas Series and apply string formatting """ @@ -190,33 +199,26 @@ def _strftime(values, date_format): return access_method(values, date_format) -class Properties: - def __init__(self, obj): - self._obj = obj +class TimeAccessor(Generic[T_DataArray]): - @staticmethod - def _tslib_field_accessor( - name: str, docstring: str = None, dtype: DTypeLike = None - ): - def f(self, dtype=dtype): - if dtype is None: - dtype = self._obj.dtype - obj_type = type(self._obj) - result = _get_date_field(self._obj.data, name, dtype) - return obj_type( - result, name=name, coords=self._obj.coords, dims=self._obj.dims - ) + __slots__ = ("_obj",) - f.__name__ = name - f.__doc__ = docstring - return property(f) + def __init__(self, obj: T_DataArray) -> None: + self._obj = obj + + def _date_field(self, name: str, dtype: DTypeLike) -> T_DataArray: + if dtype is None: + dtype = self._obj.dtype + obj_type = type(self._obj) + result = _get_date_field(self._obj.data, name, dtype) + return obj_type(result, name=name, coords=self._obj.coords, dims=self._obj.dims) - def _tslib_round_accessor(self, name, freq): + def _tslib_round_accessor(self, name: str, freq: str) -> T_DataArray: obj_type = type(self._obj) result = _round_field(self._obj.data, name, freq) return obj_type(result, name=name, coords=self._obj.coords, dims=self._obj.dims) - def floor(self, freq): + def floor(self, freq: str) -> T_DataArray: """ Round timestamps downward to specified frequency resolution. @@ -233,7 +235,7 @@ def floor(self, freq): return self._tslib_round_accessor("floor", freq) - def ceil(self, freq): + def ceil(self, freq: str) -> T_DataArray: """ Round timestamps upward to specified frequency resolution. @@ -249,7 +251,7 @@ def ceil(self, freq): """ return self._tslib_round_accessor("ceil", freq) - def round(self, freq): + def round(self, freq: str) -> T_DataArray: """ Round timestamps to specified frequency resolution. @@ -266,7 +268,7 @@ def round(self, freq): return self._tslib_round_accessor("round", freq) -class DatetimeAccessor(Properties): +class DatetimeAccessor(TimeAccessor[T_DataArray]): """Access datetime fields for DataArrays with datetime-like dtypes. Fields can be accessed through the `.dt` attribute @@ -301,7 +303,7 @@ class DatetimeAccessor(Properties): """ - def strftime(self, date_format): + def strftime(self, date_format: str) -> T_DataArray: """ Return an array of formatted strings specified by date_format, which supports the same string format as the python standard library. Details @@ -334,7 +336,7 @@ def strftime(self, date_format): result, name="strftime", coords=self._obj.coords, dims=self._obj.dims ) - def isocalendar(self): + def isocalendar(self) -> Dataset: """Dataset containing ISO year, week number, and weekday. Notes @@ -358,31 +360,48 @@ def isocalendar(self): return Dataset(data_vars) - year = Properties._tslib_field_accessor( - "year", "The year of the datetime", np.int64 - ) - month = Properties._tslib_field_accessor( - "month", "The month as January=1, December=12", np.int64 - ) - day = Properties._tslib_field_accessor("day", "The days of the datetime", np.int64) - hour = Properties._tslib_field_accessor( - "hour", "The hours of the datetime", np.int64 - ) - minute = Properties._tslib_field_accessor( - "minute", "The minutes of the datetime", np.int64 - ) - second = Properties._tslib_field_accessor( - "second", "The seconds of the datetime", np.int64 - ) - microsecond = Properties._tslib_field_accessor( - "microsecond", "The microseconds of the datetime", np.int64 - ) - nanosecond = Properties._tslib_field_accessor( - "nanosecond", "The nanoseconds of the datetime", np.int64 - ) + @property + def year(self) -> T_DataArray: + """The year of the datetime""" + return self._date_field("year", np.int64) + + @property + def month(self) -> T_DataArray: + """The month as January=1, December=12""" + return self._date_field("month", np.int64) + + @property + def day(self) -> T_DataArray: + """The days of the datetime""" + return self._date_field("day", np.int64) + + @property + def hour(self) -> T_DataArray: + """The hours of the datetime""" + return self._date_field("hour", np.int64) @property - def weekofyear(self): + def minute(self) -> T_DataArray: + """The minutes of the datetime""" + return self._date_field("minute", np.int64) + + @property + def second(self) -> T_DataArray: + """The seconds of the datetime""" + return self._date_field("second", np.int64) + + @property + def microsecond(self) -> T_DataArray: + """The microseconds of the datetime""" + return self._date_field("microsecond", np.int64) + + @property + def nanosecond(self) -> T_DataArray: + """The nanoseconds of the datetime""" + return self._date_field("nanosecond", np.int64) + + @property + def weekofyear(self) -> DataArray: "The week ordinal of the year" warnings.warn( @@ -396,64 +415,88 @@ def weekofyear(self): return weekofyear week = weekofyear - dayofweek = Properties._tslib_field_accessor( - "dayofweek", "The day of the week with Monday=0, Sunday=6", np.int64 - ) + + @property + def dayofweek(self) -> T_DataArray: + """The day of the week with Monday=0, Sunday=6""" + return self._date_field("dayofweek", np.int64) + weekday = dayofweek - weekday_name = Properties._tslib_field_accessor( - "weekday_name", "The name of day in a week", object - ) - - dayofyear = Properties._tslib_field_accessor( - "dayofyear", "The ordinal day of the year", np.int64 - ) - quarter = Properties._tslib_field_accessor("quarter", "The quarter of the date") - days_in_month = Properties._tslib_field_accessor( - "days_in_month", "The number of days in the month", np.int64 - ) + @property + def weekday_name(self) -> T_DataArray: + """The name of day in a week""" + return self._date_field("weekday_name", object) + + @property + def dayofyear(self) -> T_DataArray: + """The ordinal day of the year""" + return self._date_field("dayofyear", np.int64) + + @property + def quarter(self) -> T_DataArray: + """The quarter of the date""" + return self._date_field("quarter", np.int64) + + @property + def days_in_month(self) -> T_DataArray: + """The number of days in the month""" + return self._date_field("days_in_month", np.int64) + daysinmonth = days_in_month - season = Properties._tslib_field_accessor("season", "Season of the year", object) - - time = Properties._tslib_field_accessor( - "time", "Timestamps corresponding to datetimes", object - ) - - date = Properties._tslib_field_accessor( - "date", "Date corresponding to datetimes", object - ) - - is_month_start = Properties._tslib_field_accessor( - "is_month_start", - "Indicates whether the date is the first day of the month.", - bool, - ) - is_month_end = Properties._tslib_field_accessor( - "is_month_end", "Indicates whether the date is the last day of the month.", bool - ) - is_quarter_start = Properties._tslib_field_accessor( - "is_quarter_start", - "Indicator for whether the date is the first day of a quarter.", - bool, - ) - is_quarter_end = Properties._tslib_field_accessor( - "is_quarter_end", - "Indicator for whether the date is the last day of a quarter.", - bool, - ) - is_year_start = Properties._tslib_field_accessor( - "is_year_start", "Indicate whether the date is the first day of a year.", bool - ) - is_year_end = Properties._tslib_field_accessor( - "is_year_end", "Indicate whether the date is the last day of the year.", bool - ) - is_leap_year = Properties._tslib_field_accessor( - "is_leap_year", "Boolean indicator if the date belongs to a leap year.", bool - ) + @property + def season(self) -> T_DataArray: + """Season of the year""" + return self._date_field("season", object) + + @property + def time(self) -> T_DataArray: + """Timestamps corresponding to datetimes""" + return self._date_field("time", object) + + @property + def date(self) -> T_DataArray: + """Date corresponding to datetimes""" + return self._date_field("date", object) + + @property + def is_month_start(self) -> T_DataArray: + """Indicate whether the date is the first day of the month""" + return self._date_field("is_month_start", bool) + + @property + def is_month_end(self) -> T_DataArray: + """Indicate whether the date is the last day of the month""" + return self._date_field("is_month_end", bool) + + @property + def is_quarter_start(self) -> T_DataArray: + """Indicate whether the date is the first day of a quarter""" + return self._date_field("is_quarter_start", bool) + + @property + def is_quarter_end(self) -> T_DataArray: + """Indicate whether the date is the last day of a quarter""" + return self._date_field("is_quarter_end", bool) @property - def calendar(self): + def is_year_start(self) -> T_DataArray: + """Indicate whether the date is the first day of a year""" + return self._date_field("is_year_start", bool) + + @property + def is_year_end(self) -> T_DataArray: + """Indicate whether the date is the last day of the year""" + return self._date_field("is_year_end", bool) + + @property + def is_leap_year(self) -> T_DataArray: + """Indicate if the date belongs to a leap year""" + return self._date_field("is_leap_year", bool) + + @property + def calendar(self) -> CFCalendar: """The name of the calendar of the dates. Only relevant for arrays of :py:class:`cftime.datetime` objects, @@ -462,7 +505,7 @@ def calendar(self): return infer_calendar_name(self._obj.data) -class TimedeltaAccessor(Properties): +class TimedeltaAccessor(TimeAccessor[T_DataArray]): """Access Timedelta fields for DataArrays with Timedelta-like dtypes. Fields can be accessed through the `.dt` attribute for applicable DataArrays. @@ -502,28 +545,31 @@ class TimedeltaAccessor(Properties): * time (time) timedelta64[ns] 1 days 00:00:00 ... 5 days 18:00:00 """ - days = Properties._tslib_field_accessor( - "days", "Number of days for each element.", np.int64 - ) - seconds = Properties._tslib_field_accessor( - "seconds", - "Number of seconds (>= 0 and less than 1 day) for each element.", - np.int64, - ) - microseconds = Properties._tslib_field_accessor( - "microseconds", - "Number of microseconds (>= 0 and less than 1 second) for each element.", - np.int64, - ) - nanoseconds = Properties._tslib_field_accessor( - "nanoseconds", - "Number of nanoseconds (>= 0 and less than 1 microsecond) for each element.", - np.int64, - ) - - -class CombinedDatetimelikeAccessor(DatetimeAccessor, TimedeltaAccessor): - def __new__(cls, obj): + @property + def days(self) -> T_DataArray: + """Number of days for each element""" + return self._date_field("days", np.int64) + + @property + def seconds(self) -> T_DataArray: + """Number of seconds (>= 0 and less than 1 day) for each element""" + return self._date_field("seconds", np.int64) + + @property + def microseconds(self) -> T_DataArray: + """Number of microseconds (>= 0 and less than 1 second) for each element""" + return self._date_field("microseconds", np.int64) + + @property + def nanoseconds(self) -> T_DataArray: + """Number of nanoseconds (>= 0 and less than 1 microsecond) for each element""" + return self._date_field("nanoseconds", np.int64) + + +class CombinedDatetimelikeAccessor( + DatetimeAccessor[T_DataArray], TimedeltaAccessor[T_DataArray] +): + def __new__(cls, obj: T_DataArray) -> CombinedDatetimelikeAccessor: # CombinedDatetimelikeAccessor isn't really instatiated. Instead # we need to choose which parent (datetime or timedelta) is # appropriate. Since we're checking the dtypes anyway, we'll just @@ -537,6 +583,6 @@ def __new__(cls, obj): ) if is_np_timedelta_like(obj.dtype): - return TimedeltaAccessor(obj) + return TimedeltaAccessor(obj) # type: ignore[return-value] else: - return DatetimeAccessor(obj) + return DatetimeAccessor(obj) # type: ignore[return-value] diff --git a/xarray/core/accessor_str.py b/xarray/core/accessor_str.py index 54c9b857a7a..7f65b3add9b 100644 --- a/xarray/core/accessor_str.py +++ b/xarray/core/accessor_str.py @@ -37,27 +37,24 @@ # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +from __future__ import annotations + import codecs import re import textwrap from functools import reduce from operator import or_ as set_union -from typing import ( - Any, - Callable, - Hashable, - Mapping, - Optional, - Pattern, - Tuple, - Type, - Union, -) +from typing import TYPE_CHECKING, Any, Callable, Generic, Hashable, Mapping, Pattern from unicodedata import normalize import numpy as np from .computation import apply_ufunc +from .npcompat import DTypeLike +from .types import T_DataArray + +if TYPE_CHECKING: + from .dataarray import DataArray _cpython_optimized_encoders = ( "utf-8", @@ -112,10 +109,10 @@ def _apply_str_ufunc( *, func: Callable, obj: Any, - dtype: Union[str, np.dtype, Type] = None, - output_core_dims: Union[list, tuple] = ((),), + dtype: DTypeLike = None, + output_core_dims: list | tuple = ((),), output_sizes: Mapping[Any, int] = None, - func_args: Tuple = (), + func_args: tuple = (), func_kwargs: Mapping = {}, ) -> Any: # TODO handling of na values ? @@ -139,7 +136,7 @@ def _apply_str_ufunc( ) -class StringAccessor: +class StringAccessor(Generic[T_DataArray]): r"""Vectorized string functions for string-like arrays. Similar to pandas, fields can be accessed through the `.str` attribute @@ -204,13 +201,10 @@ class StringAccessor: __slots__ = ("_obj",) - def __init__(self, obj): + def __init__(self, obj: T_DataArray) -> None: self._obj = obj - def _stringify( - self, - invar: Any, - ) -> Union[str, bytes, Any]: + def _stringify(self, invar: Any) -> str | bytes | Any: """ Convert a string-like to the correct string/bytes type. @@ -225,10 +219,10 @@ def _apply( self, *, func: Callable, - dtype: Union[str, np.dtype, Type] = None, - output_core_dims: Union[list, tuple] = ((),), + dtype: DTypeLike = None, + output_core_dims: list | tuple = ((),), output_sizes: Mapping[Any, int] = None, - func_args: Tuple = (), + func_args: tuple = (), func_kwargs: Mapping = {}, ) -> Any: return _apply_str_ufunc( @@ -244,10 +238,10 @@ def _apply( def _re_compile( self, *, - pat: Union[str, bytes, Pattern, Any], + pat: str | bytes | Pattern | Any, flags: int = 0, case: bool = None, - ) -> Union[Pattern, Any]: + ) -> Pattern | Any: is_compiled_re = isinstance(pat, re.Pattern) if is_compiled_re and flags != 0: @@ -281,7 +275,7 @@ def func(x): else: return _apply_str_ufunc(func=func, obj=pat, dtype=np.object_) - def len(self) -> Any: + def len(self) -> T_DataArray: """ Compute the length of each string in the array. @@ -293,22 +287,19 @@ def len(self) -> Any: def __getitem__( self, - key: Union[int, slice], + key: int | slice, ) -> Any: if isinstance(key, slice): return self.slice(start=key.start, stop=key.stop, step=key.step) else: return self.get(key) - def __add__( - self, - other: Any, - ) -> Any: + def __add__(self, other: Any) -> T_DataArray: return self.cat(other, sep="") def __mul__( self, - num: Union[int, Any], + num: int | Any, ) -> Any: return self.repeat(num) @@ -327,8 +318,8 @@ def __mod__( def get( self, - i: Union[int, Any], - default: Union[str, bytes] = "", + i: int | Any, + default: str | bytes = "", ) -> Any: """ Extract character number `i` from each string in the array. @@ -360,9 +351,9 @@ def f(x, iind): def slice( self, - start: Union[int, Any] = None, - stop: Union[int, Any] = None, - step: Union[int, Any] = None, + start: int | Any = None, + stop: int | Any = None, + step: int | Any = None, ) -> Any: """ Slice substrings from each string in the array. @@ -391,9 +382,9 @@ def slice( def slice_replace( self, - start: Union[int, Any] = None, - stop: Union[int, Any] = None, - repl: Union[str, bytes, Any] = "", + start: int | Any = None, + stop: int | Any = None, + repl: str | bytes | Any = "", ) -> Any: """ Replace a positional slice of a string with another value. @@ -436,11 +427,7 @@ def func(x, istart, istop, irepl): return self._apply(func=func, func_args=(start, stop, repl)) - def cat( - self, - *others, - sep: Union[str, bytes, Any] = "", - ) -> Any: + def cat(self, *others, sep: str | bytes | Any = "") -> T_DataArray: """ Concatenate strings elementwise in the DataArray with other strings. @@ -524,8 +511,8 @@ def cat( def join( self, dim: Hashable = None, - sep: Union[str, bytes, Any] = "", - ) -> Any: + sep: str | bytes | Any = "", + ) -> T_DataArray: """ Concatenate strings in a DataArray along a particular dimension. @@ -596,7 +583,7 @@ def format( self, *args: Any, **kwargs: Any, - ) -> Any: + ) -> T_DataArray: """ Perform python string formatting on each element of the DataArray. @@ -676,7 +663,7 @@ def format( ) return self._apply(func=func, func_args=args, func_kwargs={"kwargs": kwargs}) - def capitalize(self) -> Any: + def capitalize(self) -> T_DataArray: """ Convert strings in the array to be capitalized. @@ -686,7 +673,7 @@ def capitalize(self) -> Any: """ return self._apply(func=lambda x: x.capitalize()) - def lower(self) -> Any: + def lower(self) -> T_DataArray: """ Convert strings in the array to lowercase. @@ -696,7 +683,7 @@ def lower(self) -> Any: """ return self._apply(func=lambda x: x.lower()) - def swapcase(self) -> Any: + def swapcase(self) -> T_DataArray: """ Convert strings in the array to be swapcased. @@ -706,7 +693,7 @@ def swapcase(self) -> Any: """ return self._apply(func=lambda x: x.swapcase()) - def title(self) -> Any: + def title(self) -> T_DataArray: """ Convert strings in the array to titlecase. @@ -716,7 +703,7 @@ def title(self) -> Any: """ return self._apply(func=lambda x: x.title()) - def upper(self) -> Any: + def upper(self) -> T_DataArray: """ Convert strings in the array to uppercase. @@ -726,7 +713,7 @@ def upper(self) -> Any: """ return self._apply(func=lambda x: x.upper()) - def casefold(self) -> Any: + def casefold(self) -> T_DataArray: """ Convert strings in the array to be casefolded. @@ -744,7 +731,7 @@ def casefold(self) -> Any: def normalize( self, form: str, - ) -> Any: + ) -> T_DataArray: """ Return the Unicode normal form for the strings in the datarray. @@ -763,7 +750,7 @@ def normalize( """ return self._apply(func=lambda x: normalize(form, x)) - def isalnum(self) -> Any: + def isalnum(self) -> T_DataArray: """ Check whether all characters in each string are alphanumeric. @@ -774,7 +761,7 @@ def isalnum(self) -> Any: """ return self._apply(func=lambda x: x.isalnum(), dtype=bool) - def isalpha(self) -> Any: + def isalpha(self) -> T_DataArray: """ Check whether all characters in each string are alphabetic. @@ -785,7 +772,7 @@ def isalpha(self) -> Any: """ return self._apply(func=lambda x: x.isalpha(), dtype=bool) - def isdecimal(self) -> Any: + def isdecimal(self) -> T_DataArray: """ Check whether all characters in each string are decimal. @@ -796,7 +783,7 @@ def isdecimal(self) -> Any: """ return self._apply(func=lambda x: x.isdecimal(), dtype=bool) - def isdigit(self) -> Any: + def isdigit(self) -> T_DataArray: """ Check whether all characters in each string are digits. @@ -807,7 +794,7 @@ def isdigit(self) -> Any: """ return self._apply(func=lambda x: x.isdigit(), dtype=bool) - def islower(self) -> Any: + def islower(self) -> T_DataArray: """ Check whether all characters in each string are lowercase. @@ -818,7 +805,7 @@ def islower(self) -> Any: """ return self._apply(func=lambda x: x.islower(), dtype=bool) - def isnumeric(self) -> Any: + def isnumeric(self) -> T_DataArray: """ Check whether all characters in each string are numeric. @@ -829,7 +816,7 @@ def isnumeric(self) -> Any: """ return self._apply(func=lambda x: x.isnumeric(), dtype=bool) - def isspace(self) -> Any: + def isspace(self) -> T_DataArray: """ Check whether all characters in each string are spaces. @@ -840,7 +827,7 @@ def isspace(self) -> Any: """ return self._apply(func=lambda x: x.isspace(), dtype=bool) - def istitle(self) -> Any: + def istitle(self) -> T_DataArray: """ Check whether all characters in each string are titlecase. @@ -851,7 +838,7 @@ def istitle(self) -> Any: """ return self._apply(func=lambda x: x.istitle(), dtype=bool) - def isupper(self) -> Any: + def isupper(self) -> T_DataArray: """ Check whether all characters in each string are uppercase. @@ -863,11 +850,8 @@ def isupper(self) -> Any: return self._apply(func=lambda x: x.isupper(), dtype=bool) def count( - self, - pat: Union[str, bytes, Pattern, Any], - flags: int = 0, - case: bool = None, - ) -> Any: + self, pat: str | bytes | Pattern | Any, flags: int = 0, case: bool = None + ) -> T_DataArray: """ Count occurrences of pattern in each string of the array. @@ -903,10 +887,7 @@ def count( func = lambda x, ipat: len(ipat.findall(x)) return self._apply(func=func, func_args=(pat,), dtype=int) - def startswith( - self, - pat: Union[str, bytes, Any], - ) -> Any: + def startswith(self, pat: str | bytes | Any) -> T_DataArray: """ Test if the start of each string in the array matches a pattern. @@ -929,10 +910,7 @@ def startswith( func = lambda x, y: x.startswith(y) return self._apply(func=func, func_args=(pat,), dtype=bool) - def endswith( - self, - pat: Union[str, bytes, Any], - ) -> Any: + def endswith(self, pat: str | bytes | Any) -> T_DataArray: """ Test if the end of each string in the array matches a pattern. @@ -957,10 +935,10 @@ def endswith( def pad( self, - width: Union[int, Any], + width: int | Any, side: str = "left", - fillchar: Union[str, bytes, Any] = " ", - ) -> Any: + fillchar: str | bytes | Any = " ", + ) -> T_DataArray: """ Pad strings in the array up to width. @@ -999,9 +977,9 @@ def _padder( self, *, func: Callable, - width: Union[int, Any], - fillchar: Union[str, bytes, Any] = " ", - ) -> Any: + width: int | Any, + fillchar: str | bytes | Any = " ", + ) -> T_DataArray: """ Wrapper function to handle padding operations """ @@ -1015,10 +993,8 @@ def overfunc(x, iwidth, ifillchar): return self._apply(func=overfunc, func_args=(width, fillchar)) def center( - self, - width: Union[int, Any], - fillchar: Union[str, bytes, Any] = " ", - ) -> Any: + self, width: int | Any, fillchar: str | bytes | Any = " " + ) -> T_DataArray: """ Pad left and right side of each string in the array. @@ -1043,9 +1019,9 @@ def center( def ljust( self, - width: Union[int, Any], - fillchar: Union[str, bytes, Any] = " ", - ) -> Any: + width: int | Any, + fillchar: str | bytes | Any = " ", + ) -> T_DataArray: """ Pad right side of each string in the array. @@ -1070,9 +1046,9 @@ def ljust( def rjust( self, - width: Union[int, Any], - fillchar: Union[str, bytes, Any] = " ", - ) -> Any: + width: int | Any, + fillchar: str | bytes | Any = " ", + ) -> T_DataArray: """ Pad left side of each string in the array. @@ -1095,7 +1071,7 @@ def rjust( func = self._obj.dtype.type.rjust return self._padder(func=func, width=width, fillchar=fillchar) - def zfill(self, width: Union[int, Any]) -> Any: + def zfill(self, width: int | Any) -> T_DataArray: """ Pad each string in the array by prepending '0' characters. @@ -1120,11 +1096,11 @@ def zfill(self, width: Union[int, Any]) -> Any: def contains( self, - pat: Union[str, bytes, Pattern, Any], + pat: str | bytes | Pattern | Any, case: bool = None, flags: int = 0, regex: bool = True, - ) -> Any: + ) -> T_DataArray: """ Test if pattern or regex is contained within each string of the array. @@ -1181,22 +1157,22 @@ def func(x, ipat): if case or case is None: func = lambda x, ipat: ipat in x elif self._obj.dtype.char == "U": - uppered = self._obj.str.casefold() - uppat = StringAccessor(pat).casefold() - return uppered.str.contains(uppat, regex=False) + uppered = self.casefold() + uppat = StringAccessor(pat).casefold() # type: ignore[type-var] # hack? + return uppered.str.contains(uppat, regex=False) # type: ignore[return-value] else: - uppered = self._obj.str.upper() - uppat = StringAccessor(pat).upper() - return uppered.str.contains(uppat, regex=False) + uppered = self.upper() + uppat = StringAccessor(pat).upper() # type: ignore[type-var] # hack? + return uppered.str.contains(uppat, regex=False) # type: ignore[return-value] return self._apply(func=func, func_args=(pat,), dtype=bool) def match( self, - pat: Union[str, bytes, Pattern, Any], + pat: str | bytes | Pattern | Any, case: bool = None, flags: int = 0, - ) -> Any: + ) -> T_DataArray: """ Determine if each string in the array matches a regular expression. @@ -1229,10 +1205,8 @@ def match( return self._apply(func=func, func_args=(pat,), dtype=bool) def strip( - self, - to_strip: Union[str, bytes, Any] = None, - side: str = "both", - ) -> Any: + self, to_strip: str | bytes | Any = None, side: str = "both" + ) -> T_DataArray: """ Remove leading and trailing characters. @@ -1269,10 +1243,7 @@ def strip( return self._apply(func=func, func_args=(to_strip,)) - def lstrip( - self, - to_strip: Union[str, bytes, Any] = None, - ) -> Any: + def lstrip(self, to_strip: str | bytes | Any = None) -> T_DataArray: """ Remove leading characters. @@ -1295,10 +1266,7 @@ def lstrip( """ return self.strip(to_strip, side="left") - def rstrip( - self, - to_strip: Union[str, bytes, Any] = None, - ) -> Any: + def rstrip(self, to_strip: str | bytes | Any = None) -> T_DataArray: """ Remove trailing characters. @@ -1321,11 +1289,7 @@ def rstrip( """ return self.strip(to_strip, side="right") - def wrap( - self, - width: Union[int, Any], - **kwargs, - ) -> Any: + def wrap(self, width: int | Any, **kwargs) -> T_DataArray: """ Wrap long strings in the array in paragraphs with length less than `width`. @@ -1348,14 +1312,11 @@ def wrap( wrapped : same type as values """ ifunc = lambda x: textwrap.TextWrapper(width=x, **kwargs) - tw = StringAccessor(width)._apply(func=ifunc, dtype=np.object_) + tw = StringAccessor(width)._apply(func=ifunc, dtype=np.object_) # type: ignore[type-var] # hack? func = lambda x, itw: "\n".join(itw.wrap(x)) return self._apply(func=func, func_args=(tw,)) - def translate( - self, - table: Mapping[Union[str, bytes], Union[str, bytes]], - ) -> Any: + def translate(self, table: Mapping[str | bytes, str | bytes]) -> T_DataArray: """ Map characters of each string through the given mapping table. @@ -1376,8 +1337,8 @@ def translate( def repeat( self, - repeats: Union[int, Any], - ) -> Any: + repeats: int | Any, + ) -> T_DataArray: """ Repeat each string in the array. @@ -1400,11 +1361,11 @@ def repeat( def find( self, - sub: Union[str, bytes, Any], - start: Union[int, Any] = 0, - end: Union[int, Any] = None, + sub: str | bytes | Any, + start: int | Any = 0, + end: int | Any = None, side: str = "left", - ) -> Any: + ) -> T_DataArray: """ Return lowest or highest indexes in each strings in the array where the substring is fully contained between [start:end]. @@ -1445,10 +1406,10 @@ def find( def rfind( self, - sub: Union[str, bytes, Any], - start: Union[int, Any] = 0, - end: Union[int, Any] = None, - ) -> Any: + sub: str | bytes | Any, + start: int | Any = 0, + end: int | Any = None, + ) -> T_DataArray: """ Return highest indexes in each strings in the array where the substring is fully contained between [start:end]. @@ -1477,11 +1438,11 @@ def rfind( def index( self, - sub: Union[str, bytes, Any], - start: Union[int, Any] = 0, - end: Union[int, Any] = None, + sub: str | bytes | Any, + start: int | Any = 0, + end: int | Any = None, side: str = "left", - ) -> Any: + ) -> T_DataArray: """ Return lowest or highest indexes in each strings where the substring is fully contained between [start:end]. This is the same as @@ -1528,10 +1489,10 @@ def index( def rindex( self, - sub: Union[str, bytes, Any], - start: Union[int, Any] = 0, - end: Union[int, Any] = None, - ) -> Any: + sub: str | bytes | Any, + start: int | Any = 0, + end: int | Any = None, + ) -> T_DataArray: """ Return highest indexes in each strings where the substring is fully contained between [start:end]. This is the same as @@ -1566,13 +1527,13 @@ def rindex( def replace( self, - pat: Union[str, bytes, Pattern, Any], - repl: Union[str, bytes, Callable, Any], - n: Union[int, Any] = -1, + pat: str | bytes | Pattern | Any, + repl: str | bytes | Callable | Any, + n: int | Any = -1, case: bool = None, flags: int = 0, regex: bool = True, - ) -> Any: + ) -> T_DataArray: """ Replace occurrences of pattern/regex in the array with some string. @@ -1639,7 +1600,7 @@ def replace( def extract( self, - pat: Union[str, bytes, Pattern, Any], + pat: str | bytes | Pattern | Any, dim: Hashable, case: bool = None, flags: int = 0, @@ -1783,7 +1744,7 @@ def _get_res_multi(val, pat): def extractall( self, - pat: Union[str, bytes, Pattern, Any], + pat: str | bytes | Pattern | Any, group_dim: Hashable, match_dim: Hashable, case: bool = None, @@ -1958,7 +1919,7 @@ def _get_res(val, ipat, imaxcount=maxcount, dtype=self._obj.dtype): def findall( self, - pat: Union[str, bytes, Pattern, Any], + pat: str | bytes | Pattern | Any, case: bool = None, flags: int = 0, ) -> Any: @@ -2053,9 +2014,9 @@ def _partitioner( self, *, func: Callable, - dim: Hashable, - sep: Optional[Union[str, bytes, Any]], - ) -> Any: + dim: Hashable | None, + sep: str | bytes | Any | None, + ) -> T_DataArray: """ Implements logic for `partition` and `rpartition`. """ @@ -2067,7 +2028,7 @@ def _partitioner( # _apply breaks on an empty array in this case if not self._obj.size: - return self._obj.copy().expand_dims({dim: 0}, axis=-1) + return self._obj.copy().expand_dims({dim: 0}, axis=-1) # type: ignore[return-value] arrfunc = lambda x, isep: np.array(func(x, isep), dtype=self._obj.dtype) @@ -2083,9 +2044,9 @@ def _partitioner( def partition( self, - dim: Optional[Hashable], - sep: Union[str, bytes, Any] = " ", - ) -> Any: + dim: Hashable | None, + sep: str | bytes | Any = " ", + ) -> T_DataArray: """ Split the strings in the DataArray at the first occurrence of separator `sep`. @@ -2103,7 +2064,7 @@ def partition( dim : hashable or None Name for the dimension to place the 3 elements in. If `None`, place the results as list elements in an object DataArray. - sep : str, default: " " + sep : str or bytes or array-like, default: " " String to split on. If array-like, it is broadcast. @@ -2121,9 +2082,9 @@ def partition( def rpartition( self, - dim: Optional[Hashable], - sep: Union[str, bytes, Any] = " ", - ) -> Any: + dim: Hashable | None, + sep: str | bytes | Any = " ", + ) -> T_DataArray: """ Split the strings in the DataArray at the last occurrence of separator `sep`. @@ -2141,7 +2102,7 @@ def rpartition( dim : hashable or None Name for the dimension to place the 3 elements in. If `None`, place the results as list elements in an object DataArray. - sep : str, default: " " + sep : str or bytes or array-like, default: " " String to split on. If array-like, it is broadcast. @@ -2163,9 +2124,9 @@ def _splitter( func: Callable, pre: bool, dim: Hashable, - sep: Optional[Union[str, bytes, Any]], + sep: str | bytes | Any | None, maxsplit: int, - ) -> Any: + ) -> DataArray: """ Implements logic for `split` and `rsplit`. """ @@ -2208,10 +2169,10 @@ def _dosplit(mystr, sep, maxsplit=maxsplit, dtype=self._obj.dtype): def split( self, - dim: Optional[Hashable], - sep: Union[str, bytes, Any] = None, + dim: Hashable | None, + sep: str | bytes | Any = None, maxsplit: int = -1, - ) -> Any: + ) -> DataArray: r""" Split strings in a DataArray around the given separator/delimiter `sep`. @@ -2324,10 +2285,10 @@ def split( def rsplit( self, - dim: Optional[Hashable], - sep: Union[str, bytes, Any] = None, - maxsplit: Union[int, Any] = -1, - ) -> Any: + dim: Hashable | None, + sep: str | bytes | Any = None, + maxsplit: int | Any = -1, + ) -> DataArray: r""" Split strings in a DataArray around the given separator/delimiter `sep`. @@ -2443,8 +2404,8 @@ def rsplit( def get_dummies( self, dim: Hashable, - sep: Union[str, bytes, Any] = "|", - ) -> Any: + sep: str | bytes | Any = "|", + ) -> DataArray: """ Return DataArray of dummy/indicator variables. @@ -2519,11 +2480,7 @@ def get_dummies( res.coords[dim] = vals return res - def decode( - self, - encoding: str, - errors: str = "strict", - ) -> Any: + def decode(self, encoding: str, errors: str = "strict") -> T_DataArray: """ Decode character string in the array using indicated encoding. @@ -2533,7 +2490,7 @@ def decode( The encoding to use. Please see the Python documentation `codecs standard encoders `_ section for a list of encodings handlers. - errors : str, optional + errors : str, default: "strict" The handler for encoding errors. Please see the Python documentation `codecs error handlers `_ for a list of error handlers. @@ -2549,11 +2506,7 @@ def decode( func = lambda x: decoder(x, errors)[0] return self._apply(func=func, dtype=np.str_) - def encode( - self, - encoding: str, - errors: str = "strict", - ) -> Any: + def encode(self, encoding: str, errors: str = "strict") -> T_DataArray: """ Encode character string in the array using indicated encoding. @@ -2563,7 +2516,7 @@ def encode( The encoding to use. Please see the Python documentation `codecs standard encoders `_ section for a list of encodings handlers. - errors : str, optional + errors : str, default: "strict" The handler for encoding errors. Please see the Python documentation `codecs error handlers `_ for a list of error handlers. diff --git a/xarray/core/alignment.py b/xarray/core/alignment.py index d201e3a613f..df8b3c24a91 100644 --- a/xarray/core/alignment.py +++ b/xarray/core/alignment.py @@ -16,6 +16,7 @@ Tuple, Type, TypeVar, + cast, ) import numpy as np @@ -30,6 +31,7 @@ if TYPE_CHECKING: from .dataarray import DataArray from .dataset import Dataset + from .types import JoinOptions, T_DataArray, T_DataArrayOrSet, T_Dataset DataAlignable = TypeVar("DataAlignable", bound=DataWithCoords) @@ -557,8 +559,8 @@ def align(self) -> None: def align( *objects: DataAlignable, - join="inner", - copy=True, + join: JoinOptions = "inner", + copy: bool = True, indexes=None, exclude=frozenset(), fill_value=dtypes.NA, @@ -590,7 +592,8 @@ def align( - "override": if indexes are of same size, rewrite indexes to be those of the first object with that dimension. Indexes for the same dimension must have the same size in all objects. - copy : bool, optional + + copy : bool, default: True If ``copy=True``, data in the return values is always copied. If ``copy=False`` and reindexing is unnecessary, or can be performed with only slice operations, then the output may share memory with the input. @@ -607,7 +610,7 @@ def align( Returns ------- - aligned : DataArray or Dataset + aligned : tuple of DataArray or Dataset Tuple of objects with the same type as `*objects` with aligned coordinates. @@ -763,8 +766,8 @@ def align( def deep_align( - objects, - join="inner", + objects: Iterable[Any], + join: JoinOptions = "inner", copy=True, indexes=None, exclude=frozenset(), @@ -834,7 +837,7 @@ def is_alignable(obj): if key is no_key: out[position] = aligned_obj else: - out[position][key] = aligned_obj + out[position][key] = aligned_obj # type: ignore[index] # maybe someone can fix this? # something went wrong: we should have replaced all sentinel values for arg in out: @@ -927,13 +930,15 @@ def _get_broadcast_dims_map_common_coords(args, exclude): for dim in arg.dims: if dim not in common_coords and dim not in exclude: dims_map[dim] = arg.sizes[dim] - if dim in arg.coords: - common_coords[dim] = arg.coords[dim].variable + if dim in arg._indexes: + common_coords.update(arg.xindexes.get_all_coords(dim)) return dims_map, common_coords -def _broadcast_helper(arg, exclude, dims_map, common_coords): +def _broadcast_helper( + arg: T_DataArrayOrSet, exclude, dims_map, common_coords +) -> T_DataArrayOrSet: from .dataarray import DataArray from .dataset import Dataset @@ -948,22 +953,25 @@ def _set_dims(var): return var.set_dims(var_dims_map) - def _broadcast_array(array): + def _broadcast_array(array: T_DataArray) -> T_DataArray: data = _set_dims(array.variable) coords = dict(array.coords) coords.update(common_coords) - return DataArray(data, coords, data.dims, name=array.name, attrs=array.attrs) + return array.__class__( + data, coords, data.dims, name=array.name, attrs=array.attrs + ) - def _broadcast_dataset(ds): + def _broadcast_dataset(ds: T_Dataset) -> T_Dataset: data_vars = {k: _set_dims(ds.variables[k]) for k in ds.data_vars} coords = dict(ds.coords) coords.update(common_coords) - return Dataset(data_vars, coords, ds.attrs) + return ds.__class__(data_vars, coords, ds.attrs) + # remove casts once https://github.com/python/mypy/issues/12800 is resolved if isinstance(arg, DataArray): - return _broadcast_array(arg) + return cast("T_DataArrayOrSet", _broadcast_array(arg)) elif isinstance(arg, Dataset): - return _broadcast_dataset(arg) + return cast("T_DataArrayOrSet", _broadcast_dataset(arg)) else: raise ValueError("all input must be Dataset or DataArray objects") diff --git a/xarray/core/combine.py b/xarray/core/combine.py index 78f016fdccd..fe4178eca61 100644 --- a/xarray/core/combine.py +++ b/xarray/core/combine.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import itertools import warnings from collections import Counter -from typing import Iterable, Sequence, Union +from typing import TYPE_CHECKING, Iterable, Literal, Sequence, Union import pandas as pd @@ -12,6 +14,9 @@ from .merge import merge from .utils import iterate_nested +if TYPE_CHECKING: + from .types import CombineAttrsOptions, CompatOptions, JoinOptions + def _infer_concat_order_from_positions(datasets): return dict(_infer_tile_ids_from_nested_list(datasets, ())) @@ -188,10 +193,10 @@ def _combine_nd( concat_dims, data_vars="all", coords="different", - compat="no_conflicts", + compat: CompatOptions = "no_conflicts", fill_value=dtypes.NA, - join="outer", - combine_attrs="drop", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "drop", ): """ Combines an N-dimensional structure of datasets into one by applying a @@ -250,10 +255,10 @@ def _combine_all_along_first_dim( dim, data_vars, coords, - compat, + compat: CompatOptions, fill_value=dtypes.NA, - join="outer", - combine_attrs="drop", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "drop", ): # Group into lines of datasets which must be combined along dim @@ -276,12 +281,12 @@ def _combine_all_along_first_dim( def _combine_1d( datasets, concat_dim, - compat="no_conflicts", + compat: CompatOptions = "no_conflicts", data_vars="all", coords="different", fill_value=dtypes.NA, - join="outer", - combine_attrs="drop", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "drop", ): """ Applies either concat or merge to 1D list of datasets depending on value @@ -336,8 +341,8 @@ def _nested_combine( coords, ids, fill_value=dtypes.NA, - join="outer", - combine_attrs="drop", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "drop", ): if len(datasets) == 0: @@ -377,15 +382,13 @@ def _nested_combine( def combine_nested( datasets: DATASET_HYPERCUBE, - concat_dim: Union[ - str, DataArray, None, Sequence[Union[str, "DataArray", pd.Index, None]] - ], + concat_dim: (str | DataArray | None | Sequence[str | DataArray | pd.Index | None]), compat: str = "no_conflicts", data_vars: str = "all", coords: str = "different", fill_value: object = dtypes.NA, - join: str = "outer", - combine_attrs: str = "drop", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "drop", ) -> Dataset: """ Explicitly combine an N-dimensional grid of datasets into one by using a @@ -603,9 +606,9 @@ def _combine_single_variable_hypercube( fill_value=dtypes.NA, data_vars="all", coords="different", - compat="no_conflicts", - join="outer", - combine_attrs="no_conflicts", + compat: CompatOptions = "no_conflicts", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "no_conflicts", ): """ Attempt to combine a list of Datasets into a hypercube using their @@ -659,15 +662,15 @@ def _combine_single_variable_hypercube( # TODO remove empty list default param after version 0.21, see PR4696 def combine_by_coords( - data_objects: Sequence[Union[Dataset, DataArray]] = [], - compat: str = "no_conflicts", - data_vars: str = "all", + data_objects: Iterable[Dataset | DataArray] = [], + compat: CompatOptions = "no_conflicts", + data_vars: Literal["all", "minimal", "different"] | list[str] = "all", coords: str = "different", fill_value: object = dtypes.NA, - join: str = "outer", - combine_attrs: str = "no_conflicts", - datasets: Sequence[Dataset] = None, -) -> Union[Dataset, DataArray]: + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "no_conflicts", + datasets: Iterable[Dataset] = None, +) -> Dataset | DataArray: """ Attempt to auto-magically combine the given datasets (or data arrays) @@ -695,7 +698,7 @@ def combine_by_coords( Parameters ---------- - data_objects : sequence of xarray.Dataset or sequence of xarray.DataArray + data_objects : Iterable of Datasets or DataArrays Data objects to combine. compat : {"identical", "equals", "broadcast_equals", "no_conflicts", "override"}, optional @@ -711,18 +714,19 @@ def combine_by_coords( must be equal. The returned dataset then contains the combination of all non-null values. - "override": skip comparing and pick variable from first dataset + data_vars : {"minimal", "different", "all" or list of str}, optional These data variables will be concatenated together: - * "minimal": Only data variables in which the dimension already + - "minimal": Only data variables in which the dimension already appears are included. - * "different": Data variables which are not equal (ignoring + - "different": Data variables which are not equal (ignoring attributes) across all datasets are also concatenated (as well as all for which dimension already appears). Beware: this option may load the data payload of data variables into memory if they are not already loaded. - * "all": All data variables will be concatenated. - * list of str: The listed data variables will be concatenated, in + - "all": All data variables will be concatenated. + - list of str: The listed data variables will be concatenated, in addition to the "minimal" data variables. If objects are DataArrays, `data_vars` must be "all". @@ -745,6 +749,7 @@ def combine_by_coords( - "override": if indexes are of same size, rewrite indexes to be those of the first object with that dimension. Indexes for the same dimension must have the same size in all objects. + combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \ "override"} or callable, default: "drop" A callable or a string indicating how to combine attrs of the objects being @@ -762,6 +767,8 @@ def combine_by_coords( If a callable, it must expect a sequence of ``attrs`` dicts and a context object as its only parameters. + datasets : Iterable of Datasets + Returns ------- combined : xarray.Dataset or xarray.DataArray diff --git a/xarray/core/common.py b/xarray/core/common.py index c33db4a62ea..0579a065855 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -13,6 +13,7 @@ Iterator, Mapping, TypeVar, + Union, overload, ) @@ -20,7 +21,7 @@ import pandas as pd from . import dtypes, duck_array_ops, formatting, formatting_html, ops -from .npcompat import DTypeLike +from .npcompat import DTypeLike, DTypeLikeSave from .options import OPTIONS, _get_keep_attrs from .pycompat import is_duck_dask_array from .rolling_exp import RollingExp @@ -158,6 +159,10 @@ def _repr_html_(self): return f"
{escape(repr(self))}
" return formatting_html.array_repr(self) + def __format__(self: Any, format_spec: str) -> str: + # we use numpy: scalars will print fine and arrays will raise + return self.values.__format__(format_spec) + def _iter(self: Any) -> Iterator[Any]: for n in range(len(self)): yield self[n] @@ -450,7 +455,7 @@ def assign_coords(self, coords=None, **coords_kwargs): Examples -------- - Convert longitude coordinates from 0-359 to -180-179: + Convert `DataArray` longitude coordinates from 0-359 to -180-179: >>> da = xr.DataArray( ... np.random.rand(4), @@ -490,6 +495,54 @@ def assign_coords(self, coords=None, **coords_kwargs): >>> _ = da.assign_coords({"lon_2": ("lon", lon_2)}) + Note the same method applies to `Dataset` objects. + + Convert `Dataset` longitude coordinates from 0-359 to -180-179: + + >>> temperature = np.linspace(20, 32, num=16).reshape(2, 2, 4) + >>> precipitation = 2 * np.identity(4).reshape(2, 2, 4) + >>> ds = xr.Dataset( + ... data_vars=dict( + ... temperature=(["x", "y", "time"], temperature), + ... precipitation=(["x", "y", "time"], precipitation), + ... ), + ... coords=dict( + ... lon=(["x", "y"], [[260.17, 260.68], [260.21, 260.77]]), + ... lat=(["x", "y"], [[42.25, 42.21], [42.63, 42.59]]), + ... time=pd.date_range("2014-09-06", periods=4), + ... reference_time=pd.Timestamp("2014-09-05"), + ... ), + ... attrs=dict(description="Weather-related data"), + ... ) + >>> ds + + Dimensions: (x: 2, y: 2, time: 4) + Coordinates: + lon (x, y) float64 260.2 260.7 260.2 260.8 + lat (x, y) float64 42.25 42.21 42.63 42.59 + * time (time) datetime64[ns] 2014-09-06 2014-09-07 ... 2014-09-09 + reference_time datetime64[ns] 2014-09-05 + Dimensions without coordinates: x, y + Data variables: + temperature (x, y, time) float64 20.0 20.8 21.6 22.4 ... 30.4 31.2 32.0 + precipitation (x, y, time) float64 2.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 2.0 + Attributes: + description: Weather-related data + >>> ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180)) + + Dimensions: (x: 2, y: 2, time: 4) + Coordinates: + lon (x, y) float64 -99.83 -99.32 -99.79 -99.23 + lat (x, y) float64 42.25 42.21 42.63 42.59 + * time (time) datetime64[ns] 2014-09-06 2014-09-07 ... 2014-09-09 + reference_time datetime64[ns] 2014-09-05 + Dimensions without coordinates: x, y + Data variables: + temperature (x, y, time) float64 20.0 20.8 21.6 22.4 ... 30.4 31.2 32.0 + precipitation (x, y, time) float64 2.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 2.0 + Attributes: + description: Weather-related data + Notes ----- Since ``coords_kwargs`` is a dictionary, the order of your arguments @@ -1197,8 +1250,7 @@ def where(self, cond, other=dtypes.NA, drop: bool = False): By default, these locations filled with NA. drop : bool, optional If True, coordinate labels that only correspond to False values of - the condition are dropped from the result. Mutually exclusive with - ``other``. + the condition are dropped from the result. Returns ------- @@ -1251,6 +1303,14 @@ def where(self, cond, other=dtypes.NA, drop: bool = False): [15., nan, nan, nan]]) Dimensions without coordinates: x, y + >>> a.where(a.x + a.y < 4, -1, drop=True) + + array([[ 0, 1, 2, 3], + [ 5, 6, 7, -1], + [10, 11, -1, -1], + [15, -1, -1, -1]]) + Dimensions without coordinates: x, y + See Also -------- numpy.where : corresponding numpy function @@ -1264,9 +1324,6 @@ def where(self, cond, other=dtypes.NA, drop: bool = False): cond = cond(self) if drop: - if other is not dtypes.NA: - raise ValueError("cannot set `other` if drop=True") - if not isinstance(cond, (Dataset, DataArray)): raise TypeError( f"cond argument is {cond!r} but must be a {Dataset!r} or {DataArray!r}" @@ -1521,26 +1578,51 @@ def __getitem__(self, value): raise NotImplementedError() +DTypeMaybeMapping = Union[DTypeLikeSave, Mapping[Any, DTypeLikeSave]] + + +@overload +def full_like( + other: DataArray, fill_value: Any, dtype: DTypeLikeSave = None +) -> DataArray: + ... + + @overload def full_like( - other: Dataset, - fill_value, - dtype: DTypeLike | Mapping[Any, DTypeLike] = None, + other: Dataset, fill_value: Any, dtype: DTypeMaybeMapping = None ) -> Dataset: ... @overload -def full_like(other: DataArray, fill_value, dtype: DTypeLike = None) -> DataArray: +def full_like( + other: Variable, fill_value: Any, dtype: DTypeLikeSave = None +) -> Variable: + ... + + +@overload +def full_like( + other: Dataset | DataArray, fill_value: Any, dtype: DTypeMaybeMapping = None +) -> Dataset | DataArray: ... @overload -def full_like(other: Variable, fill_value, dtype: DTypeLike = None) -> Variable: +def full_like( + other: Dataset | DataArray | Variable, + fill_value: Any, + dtype: DTypeMaybeMapping = None, +) -> Dataset | DataArray | Variable: ... -def full_like(other, fill_value, dtype=None): +def full_like( + other: Dataset | DataArray | Variable, + fill_value: Any, + dtype: DTypeMaybeMapping = None, +) -> Dataset | DataArray | Variable: """Return a new object with the same shape and type as a given object. Parameters @@ -1655,26 +1737,26 @@ def full_like(other, fill_value, dtype=None): f"fill_value must be scalar or, for datasets, a dict-like. Received {fill_value} instead." ) - if not isinstance(other, Dataset) and isinstance(dtype, Mapping): - raise ValueError( - "'dtype' cannot be dict-like when passing a DataArray or Variable" - ) - if isinstance(other, Dataset): if not isinstance(fill_value, dict): fill_value = {k: fill_value for k in other.data_vars.keys()} + dtype_: Mapping[Any, DTypeLikeSave] if not isinstance(dtype, Mapping): dtype_ = {k: dtype for k in other.data_vars.keys()} else: dtype_ = dtype data_vars = { - k: _full_like_variable(v, fill_value.get(k, dtypes.NA), dtype_.get(k, None)) + k: _full_like_variable( + v.variable, fill_value.get(k, dtypes.NA), dtype_.get(k, None) + ) for k, v in other.data_vars.items() } return Dataset(data_vars, coords=other.coords, attrs=other.attrs) elif isinstance(other, DataArray): + if isinstance(dtype, Mapping): + raise ValueError("'dtype' cannot be dict-like when passing a DataArray") return DataArray( _full_like_variable(other.variable, fill_value, dtype), dims=other.dims, @@ -1683,12 +1765,16 @@ def full_like(other, fill_value, dtype=None): name=other.name, ) elif isinstance(other, Variable): + if isinstance(dtype, Mapping): + raise ValueError("'dtype' cannot be dict-like when passing a Variable") return _full_like_variable(other, fill_value, dtype) else: raise TypeError("Expected DataArray, Dataset, or Variable") -def _full_like_variable(other, fill_value, dtype: DTypeLike = None): +def _full_like_variable( + other: Variable, fill_value: Any, dtype: DTypeLike = None +) -> Variable: """Inner function of full_like, where other must be a variable""" from .variable import Variable @@ -1709,7 +1795,38 @@ def _full_like_variable(other, fill_value, dtype: DTypeLike = None): return Variable(dims=other.dims, data=data, attrs=other.attrs) -def zeros_like(other, dtype: DTypeLike = None): +@overload +def zeros_like(other: DataArray, dtype: DTypeLikeSave = None) -> DataArray: + ... + + +@overload +def zeros_like(other: Dataset, dtype: DTypeMaybeMapping = None) -> Dataset: + ... + + +@overload +def zeros_like(other: Variable, dtype: DTypeLikeSave = None) -> Variable: + ... + + +@overload +def zeros_like( + other: Dataset | DataArray, dtype: DTypeMaybeMapping = None +) -> Dataset | DataArray: + ... + + +@overload +def zeros_like( + other: Dataset | DataArray | Variable, dtype: DTypeMaybeMapping = None +) -> Dataset | DataArray | Variable: + ... + + +def zeros_like( + other: Dataset | DataArray | Variable, dtype: DTypeMaybeMapping = None +) -> Dataset | DataArray | Variable: """Return a new object of zeros with the same shape and type as a given dataarray or dataset. @@ -1765,7 +1882,38 @@ def zeros_like(other, dtype: DTypeLike = None): return full_like(other, 0, dtype) -def ones_like(other, dtype: DTypeLike = None): +@overload +def ones_like(other: DataArray, dtype: DTypeLikeSave = None) -> DataArray: + ... + + +@overload +def ones_like(other: Dataset, dtype: DTypeMaybeMapping = None) -> Dataset: + ... + + +@overload +def ones_like(other: Variable, dtype: DTypeLikeSave = None) -> Variable: + ... + + +@overload +def ones_like( + other: Dataset | DataArray, dtype: DTypeMaybeMapping = None +) -> Dataset | DataArray: + ... + + +@overload +def ones_like( + other: Dataset | DataArray | Variable, dtype: DTypeMaybeMapping = None +) -> Dataset | DataArray | Variable: + ... + + +def ones_like( + other: Dataset | DataArray | Variable, dtype: DTypeMaybeMapping = None +) -> Dataset | DataArray | Variable: """Return a new object of ones with the same shape and type as a given dataarray or dataset. @@ -1858,7 +2006,10 @@ def _contains_cftime_datetimes(array) -> bool: def contains_cftime_datetimes(var) -> bool: """Check if an xarray.Variable contains cftime.datetime objects""" - return _contains_cftime_datetimes(var.data) + if var.dtype == np.dtype("O") and var.size > 0: + return _contains_cftime_datetimes(var.data) + else: + return False def _contains_datetime_like_objects(var) -> bool: diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 515329110af..ba1eb088bfb 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -17,24 +17,28 @@ Iterable, Mapping, Sequence, + overload, ) import numpy as np from . import dtypes, duck_array_ops, utils from .alignment import align, deep_align +from .common import zeros_like +from .duck_array_ops import datetime_to_numeric from .indexes import Index, filter_indexes_from_coords from .merge import merge_attrs, merge_coordinates_without_align from .options import OPTIONS, _get_keep_attrs from .pycompat import is_duck_dask_array -from .utils import is_dict_like +from .types import T_DataArray +from .utils import is_dict_like, is_scalar from .variable import Variable if TYPE_CHECKING: from .coordinates import Coordinates from .dataarray import DataArray from .dataset import Dataset - from .types import T_Xarray + from .types import CombineAttrsOptions, JoinOptions, T_Xarray _NO_FILL_VALUE = utils.ReprObject("") _DEFAULT_NAME = utils.ReprObject("") @@ -181,7 +185,7 @@ def _enumerate(dim): return str(alt_signature) -def result_name(objects: list) -> Any: +def result_name(objects: Iterable[Any]) -> Any: # use the same naming heuristics as pandas: # https://github.com/blaze/blaze/issues/458#issuecomment-51936356 names = {getattr(obj, "name", _DEFAULT_NAME) for obj in objects} @@ -193,7 +197,7 @@ def result_name(objects: list) -> Any: return name -def _get_coords_list(args) -> list[Coordinates]: +def _get_coords_list(args: Iterable[Any]) -> list[Coordinates]: coords_list = [] for arg in args: try: @@ -206,16 +210,16 @@ def _get_coords_list(args) -> list[Coordinates]: def build_output_coords_and_indexes( - args: list, + args: Iterable[Any], signature: _UFuncSignature, exclude_dims: AbstractSet = frozenset(), - combine_attrs: str = "override", + combine_attrs: CombineAttrsOptions = "override", ) -> tuple[list[dict[Any, Variable]], list[dict[Any, Index]]]: """Build output coordinates and indexes for an operation. Parameters ---------- - args : list + args : Iterable List of raw operation arguments. Any valid types for xarray operations are OK, e.g., scalars, Variable, DataArray, Dataset. signature : _UfuncSignature @@ -223,6 +227,22 @@ def build_output_coords_and_indexes( exclude_dims : set, optional Dimensions excluded from the operation. Coordinates along these dimensions are dropped. + combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \ + "override"} or callable, default: "drop" + A callable or a string indicating how to combine attrs of the objects being + merged: + + - "drop": empty attrs on returned Dataset. + - "identical": all attrs must be the same on every object. + - "no_conflicts": attrs from all objects are combined, any that have + the same name must also have the same value. + - "drop_conflicts": attrs from all objects are combined, any that have + the same name but different values are dropped. + - "override": skip comparing and copy attrs from the first dataset to + the result. + + If a callable, it must expect a sequence of ``attrs`` dicts and a context object + as its only parameters. Returns ------- @@ -263,11 +283,11 @@ def build_output_coords_and_indexes( def apply_dataarray_vfunc( func, *args, - signature, - join="inner", + signature: _UFuncSignature, + join: JoinOptions = "inner", exclude_dims=frozenset(), keep_attrs="override", -): +) -> tuple[DataArray, ...] | DataArray: """Apply a variable level function over DataArray, Variable and/or ndarray objects. """ @@ -292,6 +312,7 @@ def apply_dataarray_vfunc( data_vars = [getattr(a, "variable", a) for a in args] result_var = func(*data_vars) + out: tuple[DataArray, ...] | DataArray if signature.num_outputs > 1: out = tuple( DataArray( @@ -385,12 +406,12 @@ def _unpack_dict_tuples( def apply_dict_of_variables_vfunc( - func, *args, signature, join="inner", fill_value=None + func, *args, signature: _UFuncSignature, join="inner", fill_value=None ): """Apply a variable level function over dicts of DataArray, DataArray, Variable and ndarray objects. """ - args = [_as_variables_or_variable(arg) for arg in args] + args = tuple(_as_variables_or_variable(arg) for arg in args) names = join_dict_keys(args, how=join) grouped_by_name = collect_dict_values(args, names, fill_value) @@ -423,13 +444,13 @@ def _fast_dataset( def apply_dataset_vfunc( func, *args, - signature, + signature: _UFuncSignature, join="inner", dataset_join="exact", fill_value=_NO_FILL_VALUE, exclude_dims=frozenset(), keep_attrs="override", -): +) -> Dataset | tuple[Dataset, ...]: """Apply a variable level function over Dataset, dict of DataArray, DataArray, Variable and/or ndarray objects. """ @@ -452,12 +473,13 @@ def apply_dataset_vfunc( list_of_coords, list_of_indexes = build_output_coords_and_indexes( args, signature, exclude_dims, combine_attrs=keep_attrs ) - args = [getattr(arg, "data_vars", arg) for arg in args] + args = tuple(getattr(arg, "data_vars", arg) for arg in args) result_vars = apply_dict_of_variables_vfunc( func, *args, signature=signature, join=dataset_join, fill_value=fill_value ) + out: Dataset | tuple[Dataset, ...] if signature.num_outputs > 1: out = tuple( _fast_dataset(*args) @@ -639,14 +661,14 @@ def _vectorize(func, signature, output_dtypes, exclude_dims): def apply_variable_ufunc( func, *args, - signature, + signature: _UFuncSignature, exclude_dims=frozenset(), dask="forbidden", output_dtypes=None, vectorize=False, keep_attrs="override", dask_gufunc_kwargs=None, -): +) -> Variable | tuple[Variable, ...]: """Apply a ndarray level function over Variable and/or ndarray objects.""" from .variable import Variable, as_compatible_data @@ -767,7 +789,7 @@ def func(*arrays): combine_attrs=keep_attrs, ) - output = [] + output: list[Variable] = [] for dims, data in zip(output_dims, result_data): data = as_compatible_data(data) if data.ndim != len(dims): @@ -828,7 +850,7 @@ def apply_ufunc( output_core_dims: Sequence[Sequence] | None = ((),), exclude_dims: AbstractSet = frozenset(), vectorize: bool = False, - join: str = "exact", + join: JoinOptions = "exact", dataset_join: str = "exact", dataset_fill_value: object = _NO_FILL_VALUE, keep_attrs: bool | str | None = None, @@ -1352,7 +1374,9 @@ def corr(da_a, da_b, dim=None): return _cov_corr(da_a, da_b, dim=dim, method="corr") -def _cov_corr(da_a, da_b, dim=None, ddof=0, method=None): +def _cov_corr( + da_a: T_DataArray, da_b: T_DataArray, dim=None, ddof=0, method=None +) -> T_DataArray: """ Internal method for xr.cov() and xr.corr() so only have to sanitize the input arrays once and we don't repeat code. @@ -1371,9 +1395,9 @@ def _cov_corr(da_a, da_b, dim=None, ddof=0, method=None): demeaned_da_b = da_b - da_b.mean(dim=dim) # 4. Compute covariance along the given dim - # N.B. `skipna=False` is required or there is a bug when computing - # auto-covariance. E.g. Try xr.cov(da,da) for - # da = xr.DataArray([[1, 2], [1, np.nan]], dims=["x", "time"]) + # + # N.B. `skipna=True` is required or auto-covariance is computed incorrectly. E.g. + # Try xr.cov(da,da) for da = xr.DataArray([[1, 2], [1, np.nan]], dims=["x", "time"]) cov = (demeaned_da_a * demeaned_da_b).sum(dim=dim, skipna=True, min_count=1) / ( valid_count ) @@ -1827,11 +1851,10 @@ def where(cond, x, y, keep_attrs=None): """ if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) - if keep_attrs is True: # keep the attributes of x, the second parameter, by default to # be consistent with the `where` method of `DataArray` and `Dataset` - keep_attrs = lambda attrs, context: attrs[1] + keep_attrs = lambda attrs, context: getattr(x, "attrs", {}) # alignment for three arguments is complicated, so don't support it yet return apply_ufunc( @@ -1846,36 +1869,126 @@ def where(cond, x, y, keep_attrs=None): ) -def polyval(coord, coeffs, degree_dim="degree"): +@overload +def polyval(coord: DataArray, coeffs: DataArray, degree_dim: Hashable) -> DataArray: + ... + + +@overload +def polyval(coord: DataArray, coeffs: Dataset, degree_dim: Hashable) -> Dataset: + ... + + +@overload +def polyval(coord: Dataset, coeffs: DataArray, degree_dim: Hashable) -> Dataset: + ... + + +@overload +def polyval(coord: Dataset, coeffs: Dataset, degree_dim: Hashable) -> Dataset: + ... + + +@overload +def polyval( + coord: Dataset | DataArray, + coeffs: Dataset | DataArray, + degree_dim: Hashable = "degree", +) -> Dataset | DataArray: + ... + + +def polyval( + coord: Dataset | DataArray, + coeffs: Dataset | DataArray, + degree_dim: Hashable = "degree", +) -> Dataset | DataArray: """Evaluate a polynomial at specific values Parameters ---------- - coord : DataArray - The 1D coordinate along which to evaluate the polynomial. - coeffs : DataArray - Coefficients of the polynomials. - degree_dim : str, default: "degree" + coord : DataArray or Dataset + Values at which to evaluate the polynomial. + coeffs : DataArray or Dataset + Coefficients of the polynomial. + degree_dim : Hashable, default: "degree" Name of the polynomial degree dimension in `coeffs`. + Returns + ------- + DataArray or Dataset + Evaluated polynomial. + See Also -------- xarray.DataArray.polyfit - numpy.polyval + numpy.polynomial.polynomial.polyval """ - from .dataarray import DataArray - from .missing import get_clean_interp_index - x = get_clean_interp_index(coord, coord.name, strict=False) + if degree_dim not in coeffs._indexes: + raise ValueError( + f"Dimension `{degree_dim}` should be a coordinate variable with labels." + ) + if not np.issubdtype(coeffs[degree_dim].dtype, int): + raise ValueError( + f"Dimension `{degree_dim}` should be of integer dtype. Received {coeffs[degree_dim].dtype} instead." + ) + max_deg = coeffs[degree_dim].max().item() + coeffs = coeffs.reindex( + {degree_dim: np.arange(max_deg + 1)}, fill_value=0, copy=False + ) + coord = _ensure_numeric(coord) + + # using Horner's method + # https://en.wikipedia.org/wiki/Horner%27s_method + res = zeros_like(coord) + coeffs.isel({degree_dim: max_deg}, drop=True) + for deg in range(max_deg - 1, -1, -1): + res *= coord + res += coeffs.isel({degree_dim: deg}, drop=True) + + return res - deg_coord = coeffs[degree_dim] - lhs = DataArray( - np.vander(x, int(deg_coord.max()) + 1), - dims=(coord.name, degree_dim), - coords={coord.name: coord, degree_dim: np.arange(deg_coord.max() + 1)[::-1]}, - ) - return (lhs * coeffs).sum(degree_dim) +def _ensure_numeric(data: Dataset | DataArray) -> Dataset | DataArray: + """Converts all datetime64 variables to float64 + + Parameters + ---------- + data : DataArray or Dataset + Variables with possible datetime dtypes. + + Returns + ------- + DataArray or Dataset + Variables with datetime64 dtypes converted to float64. + """ + from .dataset import Dataset + + def _cfoffset(x: DataArray) -> Any: + scalar = x.compute().data[0] + if not is_scalar(scalar): + # we do not get a scalar back on dask == 2021.04.1 + scalar = scalar.item() + return type(scalar)(1970, 1, 1) + + def to_floatable(x: DataArray) -> DataArray: + if x.dtype.kind in "MO": + # datetimes (CFIndexes are object type) + offset = ( + np.datetime64("1970-01-01") if x.dtype.kind == "M" else _cfoffset(x) + ) + return x.copy( + data=datetime_to_numeric(x.data, offset=offset, datetime_unit="ns"), + ) + elif x.dtype.kind == "m": + # timedeltas + return x.astype(float) + return x + + if isinstance(data, Dataset): + return data.map(to_floatable) + else: + return to_floatable(data) def _calc_idxminmax( diff --git a/xarray/core/concat.py b/xarray/core/concat.py index f4af256430c..92e81dca4e3 100644 --- a/xarray/core/concat.py +++ b/xarray/core/concat.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Hashable, Iterable, Literal, overload +from typing import TYPE_CHECKING, Any, Hashable, Iterable, overload import pandas as pd @@ -20,24 +20,20 @@ if TYPE_CHECKING: from .dataarray import DataArray from .dataset import Dataset - -compat_options = Literal[ - "identical", "equals", "broadcast_equals", "no_conflicts", "override" -] -concat_options = Literal["all", "minimal", "different"] + from .types import CombineAttrsOptions, CompatOptions, ConcatOptions, JoinOptions @overload def concat( objs: Iterable[Dataset], dim: Hashable | DataArray | pd.Index, - data_vars: concat_options | list[Hashable] = "all", - coords: concat_options | list[Hashable] = "different", - compat: compat_options = "equals", + data_vars: ConcatOptions | list[Hashable] = "all", + coords: ConcatOptions | list[Hashable] = "different", + compat: CompatOptions = "equals", positions: Iterable[Iterable[int]] | None = None, fill_value: object = dtypes.NA, - join: str = "outer", - combine_attrs: str = "override", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "override", ) -> Dataset: ... @@ -46,13 +42,13 @@ def concat( def concat( objs: Iterable[DataArray], dim: Hashable | DataArray | pd.Index, - data_vars: concat_options | list[Hashable] = "all", - coords: concat_options | list[Hashable] = "different", - compat: compat_options = "equals", + data_vars: ConcatOptions | list[Hashable] = "all", + coords: ConcatOptions | list[Hashable] = "different", + compat: CompatOptions = "equals", positions: Iterable[Iterable[int]] | None = None, fill_value: object = dtypes.NA, - join: str = "outer", - combine_attrs: str = "override", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "override", ) -> DataArray: ... @@ -62,11 +58,11 @@ def concat( dim, data_vars="all", coords="different", - compat="equals", + compat: CompatOptions = "equals", positions=None, fill_value=dtypes.NA, - join="outer", - combine_attrs="override", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "override", ): """Concatenate xarray objects along a new or existing dimension. @@ -233,17 +229,34 @@ def concat( ) if isinstance(first_obj, DataArray): - f = _dataarray_concat + return _dataarray_concat( + objs, + dim=dim, + data_vars=data_vars, + coords=coords, + compat=compat, + positions=positions, + fill_value=fill_value, + join=join, + combine_attrs=combine_attrs, + ) elif isinstance(first_obj, Dataset): - f = _dataset_concat + return _dataset_concat( + objs, + dim=dim, + data_vars=data_vars, + coords=coords, + compat=compat, + positions=positions, + fill_value=fill_value, + join=join, + combine_attrs=combine_attrs, + ) else: raise TypeError( "can only concatenate xarray Dataset and DataArray " f"objects, got {type(first_obj)}" ) - return f( - objs, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs - ) def _calc_concat_dim_index( @@ -420,11 +433,11 @@ def _dataset_concat( dim: str | DataArray | pd.Index, data_vars: str | list[str], coords: str | list[str], - compat: str, + compat: CompatOptions, positions: Iterable[Iterable[int]] | None, fill_value: object = dtypes.NA, - join: str = "outer", - combine_attrs: str = "override", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "override", ) -> Dataset: """ Concatenate a sequence of datasets along a new or existing dimension @@ -532,7 +545,8 @@ def get_indexes(name): elif name == dim: var = ds._variables[name] if not var.dims: - yield PandasIndex([var.values.item()], dim, coord_dtype=var.dtype) + data = var.set_dims(dim).values + yield PandasIndex(data, dim, coord_dtype=var.dtype) # stack up each variable and/or index to fill-out the dataset (in order) # n.b. this loop preserves variable order, needed for groupby. @@ -608,11 +622,11 @@ def _dataarray_concat( dim: str | DataArray | pd.Index, data_vars: str | list[str], coords: str | list[str], - compat: str, + compat: CompatOptions, positions: Iterable[Iterable[int]] | None, fill_value: object = dtypes.NA, - join: str = "outer", - combine_attrs: str = "override", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "override", ) -> DataArray: from .dataarray import DataArray diff --git a/xarray/core/dask_array_compat.py b/xarray/core/dask_array_compat.py index 0e0229cc3ca..e114c238b72 100644 --- a/xarray/core/dask_array_compat.py +++ b/xarray/core/dask_array_compat.py @@ -1,14 +1,11 @@ import warnings import numpy as np -from packaging.version import Version - -from .pycompat import dask_version try: import dask.array as da except ImportError: - da = None + da = None # type: ignore def _validate_pad_output_shape(input_shape, pad_width, output_shape): @@ -57,127 +54,7 @@ def pad(array, pad_width, mode="constant", **kwargs): return padded -if dask_version > Version("2.30.0"): - ensure_minimum_chunksize = da.overlap.ensure_minimum_chunksize -else: - - # copied from dask - def ensure_minimum_chunksize(size, chunks): - """Determine new chunks to ensure that every chunk >= size - - Parameters - ---------- - size : int - The maximum size of any chunk. - chunks : tuple - Chunks along one axis, e.g. ``(3, 3, 2)`` - - Examples - -------- - >>> ensure_minimum_chunksize(10, (20, 20, 1)) - (20, 11, 10) - >>> ensure_minimum_chunksize(3, (1, 1, 3)) - (5,) - - See Also - -------- - overlap - """ - if size <= min(chunks): - return chunks - - # add too-small chunks to chunks before them - output = [] - new = 0 - for c in chunks: - if c < size: - if new > size + (size - c): - output.append(new - (size - c)) - new = size - else: - new += c - if new >= size: - output.append(new) - new = 0 - if c >= size: - new += c - if new >= size: - output.append(new) - elif len(output) >= 1: - output[-1] += new - else: - raise ValueError( - f"The overlapping depth {size} is larger than your " - f"array {sum(chunks)}." - ) - - return tuple(output) - - -if dask_version > Version("2021.03.0"): +if da is not None: sliding_window_view = da.lib.stride_tricks.sliding_window_view else: - - def sliding_window_view(x, window_shape, axis=None): - from dask.array.overlap import map_overlap - from numpy.core.numeric import normalize_axis_tuple - - from .npcompat import sliding_window_view as _np_sliding_window_view - - window_shape = ( - tuple(window_shape) if np.iterable(window_shape) else (window_shape,) - ) - - window_shape_array = np.array(window_shape) - if np.any(window_shape_array <= 0): - raise ValueError("`window_shape` must contain positive values") - - if axis is None: - axis = tuple(range(x.ndim)) - if len(window_shape) != len(axis): - raise ValueError( - f"Since axis is `None`, must provide " - f"window_shape for all dimensions of `x`; " - f"got {len(window_shape)} window_shape elements " - f"and `x.ndim` is {x.ndim}." - ) - else: - axis = normalize_axis_tuple(axis, x.ndim, allow_duplicate=True) - if len(window_shape) != len(axis): - raise ValueError( - f"Must provide matching length window_shape and " - f"axis; got {len(window_shape)} window_shape " - f"elements and {len(axis)} axes elements." - ) - - depths = [0] * x.ndim - for ax, window in zip(axis, window_shape): - depths[ax] += window - 1 - - # Ensure that each chunk is big enough to leave at least a size-1 chunk - # after windowing (this is only really necessary for the last chunk). - safe_chunks = tuple( - ensure_minimum_chunksize(d + 1, c) for d, c in zip(depths, x.chunks) - ) - x = x.rechunk(safe_chunks) - - # result.shape = x_shape_trimmed + window_shape, - # where x_shape_trimmed is x.shape with every entry - # reduced by one less than the corresponding window size. - # trim chunks to match x_shape_trimmed - newchunks = tuple( - c[:-1] + (c[-1] - d,) for d, c in zip(depths, x.chunks) - ) + tuple((window,) for window in window_shape) - - kwargs = dict( - depth=tuple((0, d) for d in depths), # Overlap on +ve side only - boundary="none", - meta=x._meta, - new_axis=range(x.ndim, x.ndim + len(axis)), - chunks=newchunks, - trim=False, - window_shape=window_shape, - axis=axis, - ) - - return map_overlap(_np_sliding_window_view, x, align_arrays=False, **kwargs) + sliding_window_view = None diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index df1e096b021..ee4516147de 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -2,6 +2,7 @@ import datetime import warnings +from os import PathLike from typing import ( TYPE_CHECKING, Any, @@ -10,8 +11,10 @@ Iterable, Literal, Mapping, + NoReturn, Sequence, cast, + overload, ) import numpy as np @@ -64,10 +67,12 @@ from .variable import IndexVariable, Variable, as_compatible_data, as_variable if TYPE_CHECKING: + from typing import TypeVar, Union + try: from dask.delayed import Delayed except ImportError: - Delayed = None + Delayed = None # type: ignore try: from cdms2 import Variable as cdms2_Variable except ImportError: @@ -77,7 +82,22 @@ except ImportError: iris_Cube = None - from .types import T_DataArray, T_Xarray + from ..backends.api import T_NetcdfEngine, T_NetcdfTypes + from .types import ( + DatetimeUnitOptions, + ErrorOptions, + ErrorOptionsWithWarn, + InterpOptions, + PadModeOptions, + PadReflectOptions, + QueryEngineOptions, + QueryParserOptions, + ReindexMethodOptions, + T_DataArray, + T_Xarray, + ) + + T_XarrayOther = TypeVar("T_XarrayOther", bound=Union["DataArray", Dataset]) def _infer_coords_and_dims( @@ -259,9 +279,9 @@ class DataArray( - mapping {coord name: (dimension name, array-like)} - mapping {coord name: (tuple of dimension names, array-like)} - dims : hashable or sequence of hashable, optional - Name(s) of the data dimension(s). Must be either a hashable - (only for 1D data) or a sequence of hashables with length equal + dims : Hashable or sequence of Hashable, optional + Name(s) of the data dimension(s). Must be either a Hashable + (only for 1D data) or a sequence of Hashables with length equal to the number of dimensions. If this argument is omitted, dimension names are taken from ``coords`` (if possible) and otherwise default to ``['dim_0', ... 'dim_n']``. @@ -353,19 +373,21 @@ class DataArray( _resample_cls = resample.DataArrayResample _weighted_cls = weighted.DataArrayWeighted - dt = utils.UncachedAccessor(CombinedDatetimelikeAccessor) + dt = utils.UncachedAccessor(CombinedDatetimelikeAccessor["DataArray"]) def __init__( self, data: Any = dtypes.NA, - coords: Sequence[tuple] | Mapping[Any, Any] | None = None, + coords: Sequence[Sequence[Any] | pd.Index | DataArray] + | Mapping[Any, Any] + | None = None, dims: Hashable | Sequence[Hashable] | None = None, name: Hashable = None, attrs: Mapping = None, # internal parameters indexes: dict[Hashable, Index] = None, fastpath: bool = False, - ): + ) -> None: if fastpath: variable = data assert dims is None @@ -416,12 +438,12 @@ def __init__( @classmethod def _construct_direct( - cls, + cls: type[T_DataArray], variable: Variable, coords: dict[Any, Variable], name: Hashable, indexes: dict[Hashable, Index], - ) -> DataArray: + ) -> T_DataArray: """Shortcut around __init__ for internal use when we want to skip costly validation """ @@ -451,8 +473,10 @@ def _replace( return type(self)(variable, coords, name=name, indexes=indexes, fastpath=True) def _replace_maybe_drop_dims( - self, variable: Variable, name: Hashable | None | Default = _default - ) -> DataArray: + self: T_DataArray, + variable: Variable, + name: Hashable | None | Default = _default, + ) -> T_DataArray: if variable.dims == self.dims and variable.shape == self.shape: coords = self._coords.copy() indexes = self._indexes @@ -474,12 +498,12 @@ def _replace_maybe_drop_dims( return self._replace(variable, coords, name, indexes=indexes) def _overwrite_indexes( - self, + self: T_DataArray, indexes: Mapping[Any, Index], coords: Mapping[Any, Variable] = None, drop_coords: list[Hashable] = None, rename_dims: Mapping[Any, Any] = None, - ) -> DataArray: + ) -> T_DataArray: """Maybe replace indexes and their corresponding coordinates.""" if not indexes: return self @@ -512,8 +536,8 @@ def _to_temp_dataset(self) -> Dataset: return self._to_dataset_whole(name=_THIS_ARRAY, shallow_copy=False) def _from_temp_dataset( - self, dataset: Dataset, name: Hashable | None | Default = _default - ) -> DataArray: + self: T_DataArray, dataset: Dataset, name: Hashable | None | Default = _default + ) -> T_DataArray: variable = dataset._variables.pop(_THIS_ARRAY) coords = dataset._variables indexes = dataset._indexes @@ -574,11 +598,11 @@ def to_dataset( Parameters ---------- - dim : hashable, optional + dim : Hashable, optional Name of the dimension on this array along which to split this array into separate variables. If not provided, this array is converted into a Dataset of one variable. - name : hashable, optional + name : Hashable, optional Name to substitute for this array's name. Only valid if ``dim`` is not provided. promote_attrs : bool, default: False @@ -723,7 +747,7 @@ def dims(self) -> tuple[Hashable, ...]: return self.variable.dims @dims.setter - def dims(self, value): + def dims(self, value: Any) -> NoReturn: raise AttributeError( "you cannot assign dims on a DataArray. Use " ".rename() or .swap_dims() instead." @@ -735,7 +759,7 @@ def _item_key_to_dict(self, key: Any) -> Mapping[Hashable, Any]: key = indexing.expanded_indexer(key, self.ndim) return dict(zip(self.dims, key)) - def _getitem_coord(self, key): + def _getitem_coord(self: T_DataArray, key: Any) -> T_DataArray: from .dataset import _get_virtual_variable try: @@ -746,7 +770,7 @@ def _getitem_coord(self, key): return self._replace_maybe_drop_dims(var, name=key) - def __getitem__(self, key: Any) -> DataArray: + def __getitem__(self: T_DataArray, key: Any) -> T_DataArray: if isinstance(key, str): return self._getitem_coord(key) else: @@ -841,19 +865,36 @@ def coords(self) -> DataArrayCoordinates: """Dictionary-like container of coordinate arrays.""" return DataArrayCoordinates(self) + @overload def reset_coords( - self, - names: Iterable[Hashable] | Hashable | None = None, + self: T_DataArray, + names: Hashable | Iterable[Hashable] | None = None, + drop: Literal[False] = False, + ) -> Dataset: + ... + + @overload + def reset_coords( + self: T_DataArray, + names: Hashable | Iterable[Hashable] | None = None, + *, + drop: Literal[True], + ) -> T_DataArray: + ... + + def reset_coords( + self: T_DataArray, + names: Hashable | Iterable[Hashable] | None = None, drop: bool = False, - ) -> None | DataArray | Dataset: + ) -> T_DataArray | Dataset: """Given names of coordinates, reset them to become variables. Parameters ---------- - names : hashable or iterable of hashable, optional + names : Hashable or iterable of Hashable, optional Name(s) of non-index coordinates in this dataset to reset into variables. By default, all non-index coordinates are reset. - drop : bool, optional + drop : bool, default: False If True, remove coordinates instead of converting them into variables. @@ -904,14 +945,14 @@ def __dask_postpersist__(self): return self._dask_finalize, (self.name, func) + args @staticmethod - def _dask_finalize(results, name, func, *args, **kwargs): + def _dask_finalize(results, name, func, *args, **kwargs) -> DataArray: ds = func(results, *args, **kwargs) variable = ds._variables.pop(_THIS_ARRAY) coords = ds._variables indexes = ds._indexes return DataArray(variable, coords, name=name, indexes=indexes, fastpath=True) - def load(self, **kwargs) -> DataArray: + def load(self: T_DataArray, **kwargs) -> T_DataArray: """Manually trigger loading of this array's data from disk or a remote source into memory and return this array. @@ -935,7 +976,7 @@ def load(self, **kwargs) -> DataArray: self._coords = new._coords return self - def compute(self, **kwargs) -> DataArray: + def compute(self: T_DataArray, **kwargs) -> T_DataArray: """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. @@ -957,7 +998,7 @@ def compute(self, **kwargs) -> DataArray: new = self.copy(deep=False) return new.load(**kwargs) - def persist(self, **kwargs) -> DataArray: + def persist(self: T_DataArray, **kwargs) -> T_DataArray: """Trigger computation in constituent dask arrays This keeps them as dask arrays but encourages them to keep data in @@ -998,7 +1039,7 @@ def copy(self: T_DataArray, deep: bool = True, data: Any = None) -> T_DataArray: Returns ------- - object : DataArray + copy : DataArray New object with dimensions, attributes, coordinates, name, encoding, and optionally data copied from original. @@ -1056,15 +1097,15 @@ def copy(self: T_DataArray, deep: bool = True, data: Any = None) -> T_DataArray: return self._replace(variable, coords, indexes=indexes) - def __copy__(self) -> DataArray: + def __copy__(self: T_DataArray) -> T_DataArray: return self.copy(deep=False) - def __deepcopy__(self, memo=None) -> DataArray: + def __deepcopy__(self: T_DataArray, memo=None) -> T_DataArray: # memo does nothing but is required for compatibility with # copy.deepcopy return self.copy(deep=True) - # mutable objects should not be hashable + # mutable objects should not be Hashable # https://github.com/python/mypy/issues/4266 __hash__ = None # type: ignore[assignment] @@ -1102,7 +1143,7 @@ def chunksizes(self) -> Mapping[Any, tuple[int, ...]]: return get_chunksizes(all_variables) def chunk( - self, + self: T_DataArray, chunks: ( int | Literal["auto"] @@ -1111,9 +1152,11 @@ def chunk( | Mapping[Any, None | int | tuple[int, ...]] ) = {}, # {} even though it's technically unsafe, is being used intentionally here (#4667) name_prefix: str = "xarray-", - token: str = None, + token: str | None = None, lock: bool = False, - ) -> DataArray: + inline_array: bool = False, + **chunks_kwargs: Any, + ) -> T_DataArray: """Coerce this array's data into a dask arrays with the given chunks. If this variable is a non-dask array, it will be converted to dask @@ -1126,7 +1169,7 @@ def chunk( Parameters ---------- - chunks : int, "auto", tuple of int or mapping of hashable to int, optional + chunks : int, "auto", tuple of int or mapping of Hashable to int, optional Chunk sizes along each dimension, e.g., ``5``, ``"auto"``, ``(5, 5)`` or ``{"x": 5, "y": 5}``. name_prefix : str, optional @@ -1136,27 +1179,57 @@ def chunk( lock : optional Passed on to :py:func:`dask.array.from_array`, if the array is not already as dask array. + inline_array: optional + Passed on to :py:func:`dask.array.from_array`, if the array is not + already as dask array. + **chunks_kwargs : {dim: chunks, ...}, optional + The keyword arguments form of ``chunks``. + One of chunks or chunks_kwargs must be provided. Returns ------- chunked : xarray.DataArray + + See Also + -------- + DataArray.chunks + DataArray.chunksizes + xarray.unify_chunks + dask.array.from_array """ - if isinstance(chunks, (tuple, list)): + if chunks is None: + warnings.warn( + "None value for 'chunks' is deprecated. " + "It will raise an error in the future. Use instead '{}'", + category=FutureWarning, + ) + chunks = {} + + if isinstance(chunks, (float, str, int)): + # ignoring type; unclear why it won't accept a Literal into the value. + chunks = dict.fromkeys(self.dims, chunks) # type: ignore + elif isinstance(chunks, (tuple, list)): chunks = dict(zip(self.dims, chunks)) + else: + chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk") ds = self._to_temp_dataset().chunk( - chunks, name_prefix=name_prefix, token=token, lock=lock + chunks, + name_prefix=name_prefix, + token=token, + lock=lock, + inline_array=inline_array, ) return self._from_temp_dataset(ds) def isel( - self, - indexers: Mapping[Any, Any] = None, + self: T_DataArray, + indexers: Mapping[Any, Any] | None = None, drop: bool = False, - missing_dims: str = "raise", + missing_dims: ErrorOptionsWithWarn = "raise", **indexers_kwargs: Any, - ) -> DataArray: - """Return a new DataArray whose data is given by integer indexing + ) -> T_DataArray: + """Return a new DataArray whose data is given by selecting indexes along the specified dimension(s). Parameters @@ -1168,7 +1241,7 @@ def isel( If DataArrays are passed as indexers, xarray-style indexing will be carried out. See :ref:`indexing` for the details. One of indexers or indexers_kwargs must be provided. - drop : bool, optional + drop : bool, default: False If ``drop=True``, drop coordinates variables indexed by integers instead of making them scalar. missing_dims : {"raise", "warn", "ignore"}, default: "raise" @@ -1180,6 +1253,10 @@ def isel( **indexers_kwargs : {dim: indexer, ...}, optional The keyword arguments form of ``indexers``. + Returns + ------- + indexed : xarray.DataArray + See Also -------- Dataset.isel @@ -1237,13 +1314,13 @@ def isel( return self._replace(variable=variable, coords=coords, indexes=indexes) def sel( - self, + self: T_DataArray, indexers: Mapping[Any, Any] = None, method: str = None, tolerance=None, drop: bool = False, **indexers_kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Return a new DataArray whose data is given by selecting index labels along the specified dimension(s). @@ -1285,10 +1362,11 @@ def sel( method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional Method to use for inexact matches: - * None (default): only exact matches - * pad / ffill: propagate last valid index value forward - * backfill / bfill: propagate next valid index value backward - * nearest: use nearest valid index value + - None (default): only exact matches + - pad / ffill: propagate last valid index value forward + - backfill / bfill: propagate next valid index value backward + - nearest: use nearest valid index value + tolerance : optional Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must @@ -1355,10 +1433,10 @@ def sel( return self._from_temp_dataset(ds) def head( - self, + self: T_DataArray, indexers: Mapping[Any, int] | int | None = None, **indexers_kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Return a new DataArray whose data is given by the the first `n` values along the specified dimension(s). Default `n` = 5 @@ -1372,10 +1450,10 @@ def head( return self._from_temp_dataset(ds) def tail( - self, + self: T_DataArray, indexers: Mapping[Any, int] | int | None = None, **indexers_kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Return a new DataArray whose data is given by the the last `n` values along the specified dimension(s). Default `n` = 5 @@ -1389,10 +1467,10 @@ def tail( return self._from_temp_dataset(ds) def thin( - self, + self: T_DataArray, indexers: Mapping[Any, int] | int | None = None, **indexers_kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Return a new DataArray whose data is given by each `n` value along the specified dimension(s). @@ -1406,8 +1484,10 @@ def thin( return self._from_temp_dataset(ds) def broadcast_like( - self, other: DataArray | Dataset, exclude: Iterable[Hashable] | None = None - ) -> DataArray: + self: T_DataArray, + other: DataArray | Dataset, + exclude: Iterable[Hashable] | None = None, + ) -> T_DataArray: """Broadcast this DataArray against another Dataset or DataArray. This is equivalent to xr.broadcast(other, self)[1] @@ -1425,7 +1505,7 @@ def broadcast_like( ---------- other : Dataset or DataArray Object against which to broadcast this array. - exclude : iterable of hashable, optional + exclude : iterable of Hashable, optional Dimensions that must not be broadcasted Returns @@ -1477,10 +1557,12 @@ def broadcast_like( dims_map, common_coords = _get_broadcast_dims_map_common_coords(args, exclude) - return _broadcast_helper(args[1], exclude, dims_map, common_coords) + return _broadcast_helper( + cast("T_DataArray", args[1]), exclude, dims_map, common_coords + ) def _reindex_callback( - self, + self: T_DataArray, aligner: alignment.Aligner, dim_pos_indexers: dict[Hashable, Any], variables: dict[Hashable, Variable], @@ -1488,7 +1570,7 @@ def _reindex_callback( fill_value: Any, exclude_dims: frozenset[Hashable], exclude_vars: frozenset[Hashable], - ) -> DataArray: + ) -> T_DataArray: """Callback called from ``Aligner`` to create a new reindexed DataArray.""" if isinstance(fill_value, dict): @@ -1511,13 +1593,13 @@ def _reindex_callback( return self._from_temp_dataset(reindexed) def reindex_like( - self, + self: T_DataArray, other: DataArray | Dataset, - method: str | None = None, + method: ReindexMethodOptions = None, tolerance: int | float | Iterable[int | float] | None = None, copy: bool = True, fill_value=dtypes.NA, - ) -> DataArray: + ) -> T_DataArray: """Conform this object onto the indexes of another object, filling in missing values with ``fill_value``. The default fill value is NaN. @@ -1534,10 +1616,11 @@ def reindex_like( Method to use for filling index values from other not found on this data array: - * None (default): don't fill gaps - * pad / ffill: propagate last valid index value forward - * backfill / bfill: propagate next valid index value backward - * nearest: use nearest valid index value + - None (default): don't fill gaps + - pad / ffill: propagate last valid index value forward + - backfill / bfill: propagate next valid index value backward + - nearest: use nearest valid index value + tolerance : optional Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must @@ -1546,7 +1629,7 @@ def reindex_like( to all values, or list-like, which applies variable tolerance per element. List-like must be the same size as the index and its dtype must exactly match the index’s type. - copy : bool, optional + copy : bool, default: True If ``copy=True``, data in the return value is always copied. If ``copy=False`` and reindexing is unnecessary, or can be performed with only slice operations, then the output may share memory with @@ -1577,14 +1660,14 @@ def reindex_like( ) def reindex( - self, + self: T_DataArray, indexers: Mapping[Any, Any] = None, - method: str = None, + method: ReindexMethodOptions = None, tolerance: int | float | Iterable[int | float] | None = None, copy: bool = True, fill_value=dtypes.NA, **indexers_kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Conform this object onto the indexes of another object, filling in missing values with ``fill_value``. The default fill value is NaN. @@ -1605,10 +1688,11 @@ def reindex( Method to use for filling index values in ``indexers`` not found on this data array: - * None (default): don't fill gaps - * pad / ffill: propagate last valid index value forward - * backfill / bfill: propagate next valid index value backward - * nearest: use nearest valid index value + - None (default): don't fill gaps + - pad / ffill: propagate last valid index value forward + - backfill / bfill: propagate next valid index value backward + - nearest: use nearest valid index value + tolerance : optional Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must @@ -1667,35 +1751,49 @@ def reindex( ) def interp( - self, - coords: Mapping[Any, Any] = None, - method: str = "linear", + self: T_DataArray, + coords: Mapping[Any, Any] | None = None, + method: InterpOptions = "linear", assume_sorted: bool = False, - kwargs: Mapping[str, Any] = None, + kwargs: Mapping[str, Any] | None = None, **coords_kwargs: Any, - ) -> DataArray: - """Multidimensional interpolation of variables. + ) -> T_DataArray: + """Interpolate a DataArray onto new coordinates + + Performs univariate or multivariate interpolation of a DataArray onto + new coordinates using scipy's interpolation routines. If interpolating + along an existing dimension, :py:class:`scipy.interpolate.interp1d` is + called. When interpolating along multiple existing dimensions, an + attempt is made to decompose the interpolation into multiple + 1-dimensional interpolations. If this is possible, + :py:class:`scipy.interpolate.interp1d` is called. Otherwise, + :py:func:`scipy.interpolate.interpn` is called. Parameters ---------- coords : dict, optional Mapping from dimension names to the new coordinates. - New coordinate can be an scalar, array-like or DataArray. + New coordinate can be a scalar, array-like or DataArray. If DataArrays are passed as new coordinates, their dimensions are used for the broadcasting. Missing values are skipped. - method : str, default: "linear" - The method used to interpolate. Choose from + method : {"linear", "nearest", "zero", "slinear", "quadratic", "cubic", "polynomial"}, default: "linear" + The method used to interpolate. The method should be supported by + the scipy interpolator: + + - ``interp1d``: {"linear", "nearest", "zero", "slinear", + "quadratic", "cubic", "polynomial"} + - ``interpn``: {"linear", "nearest"} - - {"linear", "nearest"} for multidimensional array, - - {"linear", "nearest", "zero", "slinear", "quadratic", "cubic"} for 1-dimensional array. - assume_sorted : bool, optional + If ``"polynomial"`` is passed, the ``order`` keyword argument must + also be provided. + assume_sorted : bool, default: False If False, values of x can be in any order and they are sorted first. If True, x has to be an array of monotonically increasing values. - kwargs : dict + kwargs : dict-like or None, default: None Additional keyword arguments passed to scipy's interpolator. Valid - options and their behavior depend on if 1-dimensional or - multi-dimensional interpolation is used. + options and their behavior depend whether ``interp1d`` or + ``interpn`` is used. **coords_kwargs : {dim: coordinate, ...}, optional The keyword arguments form of ``coords``. One of coords or coords_kwargs must be provided. @@ -1797,27 +1895,39 @@ def interp( return self._from_temp_dataset(ds) def interp_like( - self, + self: T_DataArray, other: DataArray | Dataset, - method: str = "linear", + method: InterpOptions = "linear", assume_sorted: bool = False, - kwargs: Mapping[str, Any] = None, - ) -> DataArray: + kwargs: Mapping[str, Any] | None = None, + ) -> T_DataArray: """Interpolate this object onto the coordinates of another object, filling out of range values with NaN. + If interpolating along a single existing dimension, + :py:class:`scipy.interpolate.interp1d` is called. When interpolating + along multiple existing dimensions, an attempt is made to decompose the + interpolation into multiple 1-dimensional interpolations. If this is + possible, :py:class:`scipy.interpolate.interp1d` is called. Otherwise, + :py:func:`scipy.interpolate.interpn` is called. + Parameters ---------- other : Dataset or DataArray Object with an 'indexes' attribute giving a mapping from dimension names to an 1d array-like, which provides coordinates upon which to index the variables in this dataset. Missing values are skipped. - method : str, default: "linear" - The method used to interpolate. Choose from + method : {"linear", "nearest", "zero", "slinear", "quadratic", "cubic", "polynomial"}, default: "linear" + The method used to interpolate. The method should be supported by + the scipy interpolator: - - {"linear", "nearest"} for multidimensional array, - - {"linear", "nearest", "zero", "slinear", "quadratic", "cubic"} for 1-dimensional array. - assume_sorted : bool, optional + - {"linear", "nearest", "zero", "slinear", "quadratic", "cubic", + "polynomial"} when ``interp1d`` is called. + - {"linear", "nearest"} when ``interpn`` is called. + + If ``"polynomial"`` is passed, the ``order`` keyword argument must + also be provided. + assume_sorted : bool, default: False If False, values of coordinates that are interpolated over can be in any order and they are sorted first. If True, interpolated coordinates are assumed to be an array of monotonically increasing @@ -1852,9 +1962,11 @@ def interp_like( ) return self._from_temp_dataset(ds) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def rename( self, - new_name_or_name_dict: Hashable | Mapping[Any, Hashable] = None, + new_name_or_name_dict: Hashable | Mapping[Any, Hashable] | None = None, **names: Hashable, ) -> DataArray: """Returns a new DataArray with renamed coordinates or a new name. @@ -1865,7 +1977,7 @@ def rename( If the argument is dict-like, it used as a mapping from old names to new names for coordinates. Otherwise, use the argument as the new name for this array. - **names : hashable, optional + **names : Hashable, optional The keyword arguments form of a mapping from old names to new names for coordinates. One of new_name_or_name_dict or names must be provided. @@ -1892,8 +2004,10 @@ def rename( return self._replace(name=new_name_or_name_dict) def swap_dims( - self, dims_dict: Mapping[Any, Hashable] = None, **dims_kwargs - ) -> DataArray: + self: T_DataArray, + dims_dict: Mapping[Any, Hashable] | None = None, + **dims_kwargs, + ) -> T_DataArray: """Returns a new DataArray with swapped dimensions. Parameters @@ -1948,10 +2062,12 @@ def swap_dims( ds = self._to_temp_dataset().swap_dims(dims_dict) return self._from_temp_dataset(ds) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def expand_dims( self, dim: None | Hashable | Sequence[Hashable] | Mapping[Any, Any] = None, - axis=None, + axis: None | int | Sequence[int] = None, **dim_kwargs: Any, ) -> DataArray: """Return a new object with an additional axis (or axes) inserted at @@ -1963,16 +2079,16 @@ def expand_dims( Parameters ---------- - dim : hashable, sequence of hashable, dict, or None, optional + dim : Hashable, sequence of Hashable, dict, or None, optional Dimensions to include on the new variable. If provided as str or sequence of str, then dimensions are inserted with length 1. If provided as a dict, then the keys are the new dimensions and the values are either integers (giving the length of the new dimensions) or sequence/ndarray (giving the coordinates of the new dimensions). - axis : int, list of int or tuple of int, or None, default: None + axis : int, sequence of int, or None, default: None Axis position(s) where new axis is to be inserted (position(s) on - the result array). If a list (or tuple) of integers is passed, + the result array). If a sequence of integers is passed, multiple axes are inserted. In this case, dim arguments should be same length list. If axis=None is passed, all the axes will be inserted to the start of the result array. @@ -1984,11 +2100,11 @@ def expand_dims( Returns ------- - expanded : same type as caller - This object, but with an additional dimension(s). + expanded : DataArray + This object, but with additional dimension(s). """ if isinstance(dim, int): - raise TypeError("dim should be hashable or sequence/mapping of hashables") + raise TypeError("dim should be Hashable or sequence/mapping of Hashables") elif isinstance(dim, Sequence) and not isinstance(dim, str): if len(dim) != len(set(dim)): raise ValueError("dims should not contain duplicate values.") @@ -2000,6 +2116,8 @@ def expand_dims( ds = self._to_temp_dataset().expand_dims(dim, axis) return self._from_temp_dataset(ds) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def set_index( self, indexes: Mapping[Any, Hashable | Sequence[Hashable]] = None, @@ -2015,9 +2133,9 @@ def set_index( Mapping from names matching dimensions and values given by (lists of) the names of existing coordinates or variables to set as new (multi-)index. - append : bool, optional + append : bool, default: False If True, append the supplied index(es) to the existing index(es). - Otherwise replace the existing index(es) (default). + Otherwise replace the existing index(es). **indexes_kwargs : optional The keyword arguments form of ``indexes``. One of indexes or indexes_kwargs must be provided. @@ -2057,6 +2175,8 @@ def set_index( ds = self._to_temp_dataset().set_index(indexes, append=append, **indexes_kwargs) return self._from_temp_dataset(ds) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def reset_index( self, dims_or_levels: Hashable | Sequence[Hashable], @@ -2066,10 +2186,10 @@ def reset_index( Parameters ---------- - dims_or_levels : hashable or sequence of hashable + dims_or_levels : Hashable or sequence of Hashable Name(s) of the dimension(s) and/or multi-index level(s) that will be reset. - drop : bool, optional + drop : bool, default: False If True, remove the specified indexes and/or multi-index levels instead of extracting them as new coordinates (default: False). @@ -2087,15 +2207,15 @@ def reset_index( return self._from_temp_dataset(ds) def reorder_levels( - self, - dim_order: Mapping[Any, Sequence[int]] = None, - **dim_order_kwargs: Sequence[int], - ) -> DataArray: + self: T_DataArray, + dim_order: Mapping[Any, Sequence[int | Hashable]] | None = None, + **dim_order_kwargs: Sequence[int | Hashable], + ) -> T_DataArray: """Rearrange index levels using input order. Parameters ---------- - dim_order : optional + dim_order dict-like of Hashable to int or Hashable: optional Mapping from names matching dimensions and values given by lists representing new level orders. Every given dimension must have a multi-index. @@ -2113,12 +2233,12 @@ def reorder_levels( return self._from_temp_dataset(ds) def stack( - self, - dimensions: Mapping[Any, Sequence[Hashable]] = None, - create_index: bool = True, + self: T_DataArray, + dimensions: Mapping[Any, Sequence[Hashable]] | None = None, + create_index: bool | None = True, index_cls: type[Index] = PandasMultiIndex, **dimensions_kwargs: Sequence[Hashable], - ) -> DataArray: + ) -> T_DataArray: """ Stack any number of existing dimensions into a single new dimension. @@ -2127,14 +2247,14 @@ def stack( Parameters ---------- - dimensions : mapping of hashable to sequence of hashable + dimensions : mapping of Hashable to sequence of Hashable Mapping of the form `new_name=(dim1, dim2, ...)`. Names of new dimensions, and the existing dimensions that they replace. An ellipsis (`...`) will be replaced by all unlisted dimensions. Passing a list containing an ellipsis (`stacked_dim=[...]`) will stack over all dimensions. - create_index : bool, optional - If True (default), create a multi-index for each of the stacked dimensions. + create_index : bool or None, default: True + If True, create a multi-index for each of the stacked dimensions. If False, don't create any index. If None, create a multi-index only if exactly one single (1-d) coordinate index is found for every dimension to stack. @@ -2185,6 +2305,8 @@ def stack( ) return self._from_temp_dataset(ds) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def unstack( self, dim: Hashable | Sequence[Hashable] | None = None, @@ -2199,16 +2321,16 @@ def unstack( Parameters ---------- - dim : hashable or sequence of hashable, optional + dim : Hashable or sequence of Hashable, optional Dimension(s) over which to unstack. By default unstacks all MultiIndexes. fill_value : scalar or dict-like, default: nan - value to be filled. If a dict-like, maps variable names to + Value to be filled. If a dict-like, maps variable names to fill values. Use the data array's name to refer to its name. If not provided or if the dict-like does not contain all variables, the dtype's NA value will be used. sparse : bool, default: False - use sparse-array if True + Use sparse-array if True Returns ------- @@ -2248,7 +2370,7 @@ def unstack( ds = self._to_temp_dataset().unstack(dim, fill_value, sparse) return self._from_temp_dataset(ds) - def to_unstacked_dataset(self, dim, level=0): + def to_unstacked_dataset(self, dim: Hashable, level: int | Hashable = 0) -> Dataset: """Unstack DataArray expanding to Dataset along a given level of a stacked coordinate. @@ -2256,9 +2378,9 @@ def to_unstacked_dataset(self, dim, level=0): Parameters ---------- - dim : str + dim : Hashable Name of existing dimension to unstack - level : int or str + level : int or Hashable, default: 0 The MultiIndex level to expand to a dataset along. Can either be the integer index of the level or its name. @@ -2314,16 +2436,16 @@ def to_unstacked_dataset(self, dim, level=0): return Dataset(data_dict) def transpose( - self, + self: T_DataArray, *dims: Hashable, transpose_coords: bool = True, - missing_dims: str = "raise", - ) -> DataArray: + missing_dims: ErrorOptionsWithWarn = "raise", + ) -> T_DataArray: """Return a new DataArray object with transposed dimensions. Parameters ---------- - *dims : hashable, optional + *dims : Hashable, optional By default, reverse the dimensions. Otherwise, reorder the dimensions to this order. transpose_coords : bool, default: True @@ -2364,20 +2486,25 @@ def transpose( return self._replace(variable) @property - def T(self) -> DataArray: + def T(self: T_DataArray) -> T_DataArray: return self.transpose() + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def drop_vars( - self, names: Hashable | Iterable[Hashable], *, errors: str = "raise" + self, + names: Hashable | Iterable[Hashable], + *, + errors: ErrorOptions = "raise", ) -> DataArray: """Returns an array with dropped variables. Parameters ---------- - names : hashable or iterable of hashable + names : Hashable or iterable of Hashable Name(s) of variables to drop. - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if any of the variable + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the variable passed are not in the dataset. If 'ignore', any given names that are in the DataArray are dropped and no error is raised. @@ -2390,13 +2517,13 @@ def drop_vars( return self._from_temp_dataset(ds) def drop( - self, - labels: Mapping = None, - dim: Hashable = None, + self: T_DataArray, + labels: Mapping[Any, Any] | None = None, + dim: Hashable | None = None, *, - errors: str = "raise", + errors: ErrorOptions = "raise", **labels_kwargs, - ) -> DataArray: + ) -> T_DataArray: """Backward compatible method based on `drop_vars` and `drop_sel` Using either `drop_vars` or `drop_sel` is encouraged @@ -2406,24 +2533,24 @@ def drop( DataArray.drop_vars DataArray.drop_sel """ - ds = self._to_temp_dataset().drop(labels, dim, errors=errors) + ds = self._to_temp_dataset().drop(labels, dim, errors=errors, **labels_kwargs) return self._from_temp_dataset(ds) def drop_sel( - self, - labels: Mapping[Any, Any] = None, + self: T_DataArray, + labels: Mapping[Any, Any] | None = None, *, - errors: str = "raise", + errors: ErrorOptions = "raise", **labels_kwargs, - ) -> DataArray: + ) -> T_DataArray: """Drop index labels from this DataArray. Parameters ---------- - labels : mapping of hashable to Any + labels : mapping of Hashable to Any Index labels to drop - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the index labels passed are not in the dataset. If 'ignore', any given labels that are in the dataset are dropped and no error is raised. @@ -2440,12 +2567,14 @@ def drop_sel( ds = self._to_temp_dataset().drop_sel(labels, errors=errors) return self._from_temp_dataset(ds) - def drop_isel(self, indexers=None, **indexers_kwargs): + def drop_isel( + self: T_DataArray, indexers: Mapping[Any, Any] | None = None, **indexers_kwargs + ) -> T_DataArray: """Drop index positions from this DataArray. Parameters ---------- - indexers : mapping of hashable to Any + indexers : mapping of Hashable to Any or None, default: None Index locations to drop **indexers_kwargs : {dim: position, ...}, optional The keyword arguments form of ``dim`` and ``positions`` @@ -2462,29 +2591,35 @@ def drop_isel(self, indexers=None, **indexers_kwargs): dataset = dataset.drop_isel(indexers=indexers, **indexers_kwargs) return self._from_temp_dataset(dataset) - def dropna(self, dim: Hashable, how: str = "any", thresh: int = None) -> DataArray: + def dropna( + self: T_DataArray, + dim: Hashable, + how: Literal["any", "all"] = "any", + thresh: int | None = None, + ) -> T_DataArray: """Returns a new array with dropped labels for missing values along the provided dimension. Parameters ---------- - dim : hashable + dim : Hashable Dimension along which to drop missing values. Dropping along multiple dimensions simultaneously is not yet supported. - how : {"any", "all"}, optional - * any : if any NA values are present, drop that label - * all : if all values are NA, drop that label - thresh : int, default: None + how : {"any", "all"}, default: "any" + - any : if any NA values are present, drop that label + - all : if all values are NA, drop that label + + thresh : int or None, default: None If supplied, require this many non-NA values. Returns ------- - DataArray + dropped : DataArray """ ds = self._to_temp_dataset().dropna(dim, how=how, thresh=thresh) return self._from_temp_dataset(ds) - def fillna(self, value: Any) -> DataArray: + def fillna(self: T_DataArray, value: Any) -> T_DataArray: """Fill missing values in this object. This operation follows the normal broadcasting and alignment rules that @@ -2501,7 +2636,7 @@ def fillna(self, value: Any) -> DataArray: Returns ------- - DataArray + filled : DataArray """ if utils.is_dict_like(value): raise TypeError( @@ -2512,27 +2647,34 @@ def fillna(self, value: Any) -> DataArray: return out def interpolate_na( - self, - dim: Hashable = None, - method: str = "linear", - limit: int = None, + self: T_DataArray, + dim: Hashable | None = None, + method: InterpOptions = "linear", + limit: int | None = None, use_coordinate: bool | str = True, max_gap: ( - int | float | str | pd.Timedelta | np.timedelta64 | datetime.timedelta + None + | int + | float + | str + | pd.Timedelta + | np.timedelta64 + | datetime.timedelta ) = None, - keep_attrs: bool = None, + keep_attrs: bool | None = None, **kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Fill in NaNs by interpolating according to different methods. Parameters ---------- dim : str Specifies the dimension along which to interpolate. - method : str, optional + method : {"linear", "nearest", "zero", "slinear", "quadratic", "cubic", "polynomial", \ + "barycentric", "krog", "pchip", "spline", "akima"}, default: "linear" String indicating which method to use for interpolation: - - 'linear': linear interpolation (Default). Additional keyword + - 'linear': linear interpolation. Additional keyword arguments are passed to :py:func:`numpy.interp` - 'nearest', 'zero', 'slinear', 'quadratic', 'cubic', 'polynomial': are passed to :py:func:`scipy.interpolate.interp1d`. If @@ -2540,13 +2682,14 @@ def interpolate_na( provided. - 'barycentric', 'krog', 'pchip', 'spline', 'akima': use their respective :py:class:`scipy.interpolate` classes. + use_coordinate : bool or str, default: True Specifies which index to use as the x values in the interpolation formulated as `y = f(x)`. If False, values are treated as if eqaully-spaced along ``dim``. If True, the IndexVariable `dim` is used. If ``use_coordinate`` is a string, it specifies the name of a coordinate variariable to use as the index. - limit : int, default: None + limit : int or None, default: None Maximum number of consecutive NaNs to fill. Must be greater than 0 or None for no limit. This filling is done regardless of the size of the gap in the data. To only interpolate over gaps less than a given length, @@ -2574,7 +2717,7 @@ def interpolate_na( * x (x) int64 0 1 2 3 4 5 6 7 8 The gap lengths are 3-0 = 3; 6-3 = 3; and 8-6 = 2 respectively - keep_attrs : bool, default: True + keep_attrs : bool or None, default: None If True, the dataarray's attributes (`attrs`) will be copied from the original object to the new one. If False, the new object will be returned without attributes. @@ -2627,17 +2770,19 @@ def interpolate_na( **kwargs, ) - def ffill(self, dim: Hashable, limit: int = None) -> DataArray: + def ffill( + self: T_DataArray, dim: Hashable, limit: int | None = None + ) -> T_DataArray: """Fill NaN values by propagating values forward *Requires bottleneck.* Parameters ---------- - dim : hashable + dim : Hashable Specifies the dimension along which to propagate values when filling. - limit : int, default: None + limit : int or None, default: None The maximum number of consecutive NaN values to forward fill. In other words, if there is a gap with more than this number of consecutive NaNs, it will only be partially filled. Must be greater @@ -2646,13 +2791,15 @@ def ffill(self, dim: Hashable, limit: int = None) -> DataArray: Returns ------- - DataArray + filled : DataArray """ from .missing import ffill return ffill(self, dim, limit=limit) - def bfill(self, dim: Hashable, limit: int = None) -> DataArray: + def bfill( + self: T_DataArray, dim: Hashable, limit: int | None = None + ) -> T_DataArray: """Fill NaN values by propagating values backward *Requires bottleneck.* @@ -2662,7 +2809,7 @@ def bfill(self, dim: Hashable, limit: int = None) -> DataArray: dim : str Specifies the dimension along which to propagate values when filling. - limit : int, default: None + limit : int or None, default: None The maximum number of consecutive NaN values to backward fill. In other words, if there is a gap with more than this number of consecutive NaNs, it will only be partially filled. Must be greater @@ -2671,13 +2818,13 @@ def bfill(self, dim: Hashable, limit: int = None) -> DataArray: Returns ------- - DataArray + filled : DataArray """ from .missing import bfill return bfill(self, dim, limit=limit) - def combine_first(self, other: DataArray) -> DataArray: + def combine_first(self: T_DataArray, other: T_DataArray) -> T_DataArray: """Combine two DataArray objects, with union of coordinates. This operation follows the normal broadcasting and alignment rules of @@ -2696,15 +2843,15 @@ def combine_first(self, other: DataArray) -> DataArray: return ops.fillna(self, other, join="outer") def reduce( - self, + self: T_DataArray, func: Callable[..., Any], dim: None | Hashable | Sequence[Hashable] = None, *, axis: None | int | Sequence[int] = None, - keep_attrs: bool = None, + keep_attrs: bool | None = None, keepdims: bool = False, **kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Reduce this array by applying `func` along some dimension(s). Parameters @@ -2713,14 +2860,14 @@ def reduce( Function which can be called in the form `f(x, axis=axis, **kwargs)` to return the result of reducing an np.ndarray over an integer valued axis. - dim : hashable or sequence of hashable, optional + dim : Hashable or sequence of Hashable, optional Dimension(s) over which to apply `func`. axis : int or sequence of int, optional Axis(es) over which to repeatedly apply `func`. Only one of the 'dim' and 'axis' arguments can be supplied. If neither are supplied, then the reduction is calculated over the flattened array (by calling `f(x)` without an axis argument). - keep_attrs : bool, optional + keep_attrs : bool or None, optional If True, the variable's attributes (`attrs`) will be copied from the original object to the new one. If False (default), the new object will be returned without attributes. @@ -2754,6 +2901,11 @@ def to_pandas(self) -> DataArray | pd.Series | pd.DataFrame: Only works for arrays with 2 or fewer dimensions. The DataArray constructor performs the inverse transformation. + + Returns + ------- + result : DataArray | Series | DataFrame + DataArray, pandas Series or pandas DataFrame. """ # TODO: consolidate the info about pandas constructors and the # attributes that correspond to their indexes into a separate module? @@ -2762,14 +2914,14 @@ def to_pandas(self) -> DataArray | pd.Series | pd.DataFrame: constructor = constructors[self.ndim] except KeyError: raise ValueError( - f"cannot convert arrays with {self.ndim} dimensions into " - "pandas objects" + f"Cannot convert arrays with {self.ndim} dimensions into " + "pandas objects. Requires 2 or fewer dimensions." ) indexes = [self.get_index(dim) for dim in self.dims] return constructor(self.values, *indexes) def to_dataframe( - self, name: Hashable = None, dim_order: list[Hashable] = None + self, name: Hashable | None = None, dim_order: Sequence[Hashable] | None = None ) -> pd.DataFrame: """Convert this array and its coordinates into a tidy pandas.DataFrame. @@ -2782,9 +2934,9 @@ def to_dataframe( Parameters ---------- - name + name: Hashable or None, optional Name to give to this array (required if unnamed). - dim_order + dim_order: Sequence of Hashable or None, optional Hierarchical dimension order for the resulting dataframe. Array content is transposed to this order and then written out as flat vectors in contiguous order, so the last dimension in this list @@ -2797,12 +2949,13 @@ def to_dataframe( Returns ------- - result + result: DataFrame DataArray as a pandas DataFrame. See also -------- DataArray.to_pandas + DataArray.to_series """ if name is None: name = self.name @@ -2836,6 +2989,16 @@ def to_series(self) -> pd.Series: The Series is indexed by the Cartesian product of index coordinates (in the form of a :py:class:`pandas.MultiIndex`). + + Returns + ------- + result : Series + DataArray as a pandas Series. + + See also + -------- + DataArray.to_pandas + DataArray.to_dataframe """ index = self.coords.to_index() return pd.Series(self.values.reshape(-1), index=index, name=self.name) @@ -2858,10 +3021,140 @@ def to_masked_array(self, copy: bool = True) -> np.ma.MaskedArray: isnull = pd.isnull(values) return np.ma.MaskedArray(data=values, mask=isnull, copy=copy) - def to_netcdf(self, *args, **kwargs) -> bytes | Delayed | None: - """Write DataArray contents to a netCDF file. + # path=None writes to bytes + @overload + def to_netcdf( + self, + path: None = None, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: bool = True, + invalid_netcdf: bool = False, + ) -> bytes: + ... - All parameters are passed directly to :py:meth:`xarray.Dataset.to_netcdf`. + # default return None + @overload + def to_netcdf( + self, + path: str | PathLike, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: Literal[True] = True, + invalid_netcdf: bool = False, + ) -> None: + ... + + # compute=False returns dask.Delayed + @overload + def to_netcdf( + self, + path: str | PathLike, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + *, + compute: Literal[False], + invalid_netcdf: bool = False, + ) -> Delayed: + ... + + def to_netcdf( + self, + path: str | PathLike | None = None, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: bool = True, + invalid_netcdf: bool = False, + ) -> bytes | Delayed | None: + """Write dataset contents to a netCDF file. + + Parameters + ---------- + path : str, path-like or None, optional + Path to which to save this dataset. File-like objects are only + supported by the scipy engine. If no path is provided, this + function returns the resulting netCDF file as bytes; in this case, + we need to use scipy, which does not support netCDF version 4 (the + default format becomes NETCDF3_64BIT). + mode : {"w", "a"}, default: "w" + Write ('w') or append ('a') mode. If mode='w', any existing file at + this location will be overwritten. If mode='a', existing variables + will be overwritten. + format : {"NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_64BIT", \ + "NETCDF3_CLASSIC"}, optional + File format for the resulting netCDF file: + + * NETCDF4: Data is stored in an HDF5 file, using netCDF4 API + features. + * NETCDF4_CLASSIC: Data is stored in an HDF5 file, using only + netCDF 3 compatible API features. + * NETCDF3_64BIT: 64-bit offset version of the netCDF 3 file format, + which fully supports 2+ GB files, but is only compatible with + clients linked against netCDF version 3.6.0 or later. + * NETCDF3_CLASSIC: The classic netCDF 3 file format. It does not + handle 2+ GB files very well. + + All formats are supported by the netCDF4-python library. + scipy.io.netcdf only supports the last two formats. + + The default format is NETCDF4 if you are saving a file to disk and + have the netCDF4-python library available. Otherwise, xarray falls + back to using scipy to write netCDF files and defaults to the + NETCDF3_64BIT format (scipy does not support netCDF4). + group : str, optional + Path to the netCDF4 group in the given file to open (only works for + format='NETCDF4'). The group(s) will be created if necessary. + engine : {"netcdf4", "scipy", "h5netcdf"}, optional + Engine to use when writing netCDF files. If not provided, the + default engine is chosen based on available dependencies, with a + preference for 'netcdf4' if writing to a file on disk. + encoding : dict, optional + Nested dictionary with variable names as keys and dictionaries of + variable specific encodings as values, e.g., + ``{"my_variable": {"dtype": "int16", "scale_factor": 0.1, + "zlib": True}, ...}`` + + The `h5netcdf` engine supports both the NetCDF4-style compression + encoding parameters ``{"zlib": True, "complevel": 9}`` and the h5py + ones ``{"compression": "gzip", "compression_opts": 9}``. + This allows using any compression plugin installed in the HDF5 + library, e.g. LZF. + + unlimited_dims : iterable of Hashable, optional + Dimension(s) that should be serialized as unlimited dimensions. + By default, no dimensions are treated as unlimited dimensions. + Note that unlimited_dims may also be set via + ``dataset.encoding["unlimited_dims"]``. + compute: bool, default: True + If true compute immediately, otherwise return a + ``dask.delayed.Delayed`` object that can be computed later. + invalid_netcdf: bool, default: False + Only valid along with ``engine="h5netcdf"``. If True, allow writing + hdf5 files which are invalid netcdf as described in + https://github.com/h5netcdf/h5netcdf. + + Returns + ------- + store: bytes or Delayed or None + * ``bytes`` if path is None + * ``dask.delayed.Delayed`` if compute is False + * None otherwise Notes ----- @@ -2875,7 +3168,7 @@ def to_netcdf(self, *args, **kwargs) -> bytes | Delayed | None: -------- Dataset.to_netcdf """ - from ..backends.api import DATAARRAY_NAME, DATAARRAY_VARIABLE + from ..backends.api import DATAARRAY_NAME, DATAARRAY_VARIABLE, to_netcdf if self.name is None: # If no name is set then use a generic xarray name @@ -2889,9 +3182,21 @@ def to_netcdf(self, *args, **kwargs) -> bytes | Delayed | None: # No problems with the name - so we're fine! dataset = self.to_dataset() - return dataset.to_netcdf(*args, **kwargs) + return to_netcdf( # type: ignore # mypy cannot resolve the overloads:( + dataset, + path, + mode=mode, + format=format, + group=group, + engine=engine, + encoding=encoding, + unlimited_dims=unlimited_dims, + compute=compute, + multifile=False, + invalid_netcdf=invalid_netcdf, + ) - def to_dict(self, data: bool = True) -> dict: + def to_dict(self, data: bool = True, encoding: bool = False) -> dict[str, Any]: """ Convert this xarray.DataArray into a dictionary following xarray naming conventions. @@ -2902,22 +3207,31 @@ def to_dict(self, data: bool = True) -> dict: Parameters ---------- - data : bool, optional + data : bool, default: True Whether to include the actual data in the dictionary. When set to False, returns just the schema. + encoding : bool, default: False + Whether to include the Dataset's encoding in the dictionary. + + Returns + ------- + dict: dict See Also -------- DataArray.from_dict + Dataset.to_dict """ d = self.variable.to_dict(data=data) d.update({"coords": {}, "name": self.name}) for k in self.coords: d["coords"][k] = self.coords[k].variable.to_dict(data=data) + if encoding: + d["encoding"] = dict(self.encoding) return d @classmethod - def from_dict(cls, d: dict) -> DataArray: + def from_dict(cls: type[T_DataArray], d: Mapping[str, Any]) -> T_DataArray: """Convert a dictionary into an xarray.DataArray Parameters @@ -2979,6 +3293,9 @@ def from_dict(cls, d: dict) -> DataArray: raise ValueError("cannot convert dict without the key 'data''") else: obj = cls(data, coords, d.get("dims"), d.get("name"), d.get("attrs")) + + obj.encoding.update(d.get("encoding", {})) + return obj @classmethod @@ -2990,12 +3307,18 @@ def from_series(cls, series: pd.Series, sparse: bool = False) -> DataArray: values with NaN). Thus this operation should be the inverse of the `to_series` method. - If sparse=True, creates a sparse array instead of a dense NumPy array. - Requires the pydata/sparse package. + Parameters + ---------- + series : Series + Pandas Series object to convert. + sparse : bool, default: False + If sparse=True, creates a sparse array instead of a dense NumPy array. + Requires the pydata/sparse package. See Also -------- - xarray.Dataset.from_dataframe + DataArray.to_series + Dataset.from_dataframe """ temp_name = "__temporary_name" df = pd.DataFrame({temp_name: series}) @@ -3030,7 +3353,7 @@ def from_iris(cls, cube: iris_Cube) -> DataArray: return from_iris(cube) - def _all_compat(self, other: DataArray, compat_str: str) -> bool: + def _all_compat(self: T_DataArray, other: T_DataArray, compat_str: str) -> bool: """Helper function for equals, broadcast_equals, and identical""" def compat(x, y): @@ -3040,11 +3363,21 @@ def compat(x, y): self, other ) - def broadcast_equals(self, other: DataArray) -> bool: + def broadcast_equals(self: T_DataArray, other: T_DataArray) -> bool: """Two DataArrays are broadcast equal if they are equal after broadcasting them against each other such that they have the same dimensions. + Parameters + ---------- + other : DataArray + DataArray to compare to. + + Returns + ---------- + equal : bool + True if the two DataArrays are broadcast equal. + See Also -------- DataArray.equals @@ -3055,7 +3388,7 @@ def broadcast_equals(self, other: DataArray) -> bool: except (TypeError, AttributeError): return False - def equals(self, other: DataArray) -> bool: + def equals(self: T_DataArray, other: T_DataArray) -> bool: """True if two DataArrays have the same dimensions, coordinates and values; otherwise False. @@ -3065,6 +3398,16 @@ def equals(self, other: DataArray) -> bool: This method is necessary because `v1 == v2` for ``DataArray`` does element-wise comparisons (like numpy.ndarrays). + Parameters + ---------- + other : DataArray + DataArray to compare to. + + Returns + ---------- + equal : bool + True if the two DataArrays are equal. + See Also -------- DataArray.broadcast_equals @@ -3075,10 +3418,20 @@ def equals(self, other: DataArray) -> bool: except (TypeError, AttributeError): return False - def identical(self, other: DataArray) -> bool: + def identical(self: T_DataArray, other: T_DataArray) -> bool: """Like equals, but also checks the array name and attributes, and attributes on all coordinates. + Parameters + ---------- + other : DataArray + DataArray to compare to. + + Returns + ---------- + equal : bool + True if the two DataArrays are identical. + See Also -------- DataArray.broadcast_equals @@ -3098,19 +3451,19 @@ def _result_name(self, other: Any = None) -> Hashable | None: else: return None - def __array_wrap__(self, obj, context=None) -> DataArray: + def __array_wrap__(self: T_DataArray, obj, context=None) -> T_DataArray: new_var = self.variable.__array_wrap__(obj, context) return self._replace(new_var) - def __matmul__(self, obj): + def __matmul__(self: T_DataArray, obj: T_DataArray) -> T_DataArray: return self.dot(obj) - def __rmatmul__(self, other): + def __rmatmul__(self: T_DataArray, other: T_DataArray) -> T_DataArray: # currently somewhat duplicative, as only other DataArrays are # compatible with matmul return computation.dot(other, self) - def _unary_op(self, f: Callable, *args, **kwargs): + def _unary_op(self: T_DataArray, f: Callable, *args, **kwargs) -> T_DataArray: keep_attrs = kwargs.pop("keep_attrs", None) if keep_attrs is None: keep_attrs = _get_keep_attrs(default=True) @@ -3126,16 +3479,16 @@ def _unary_op(self, f: Callable, *args, **kwargs): return da def _binary_op( - self, - other, + self: T_DataArray, + other: Any, f: Callable, reflexive: bool = False, - ): + ) -> T_DataArray: if isinstance(other, (Dataset, groupby.GroupBy)): return NotImplemented if isinstance(other, DataArray): align_type = OPTIONS["arithmetic_join"] - self, other = align(self, other, join=align_type, copy=False) + self, other = align(self, other, join=align_type, copy=False) # type: ignore other_variable = getattr(other, "variable", other) other_coords = getattr(other, "coords", None) @@ -3149,7 +3502,7 @@ def _binary_op( return self._replace(variable, coords, name, indexes=indexes) - def _inplace_binary_op(self, other, f: Callable): + def _inplace_binary_op(self: T_DataArray, other: Any, f: Callable) -> T_DataArray: if isinstance(other, groupby.GroupBy): raise TypeError( "in-place operations between a DataArray and " @@ -3210,16 +3563,18 @@ def _title_for_slice(self, truncate: int = 50) -> str: return title - def diff(self, dim: Hashable, n: int = 1, label: Hashable = "upper") -> DataArray: + def diff( + self: T_DataArray, dim: Hashable, n: int = 1, label: Hashable = "upper" + ) -> T_DataArray: """Calculate the n-th order discrete difference along given axis. Parameters ---------- - dim : hashable + dim : Hashable Dimension over which to calculate the finite difference. - n : int, optional + n : int, default: 1 The number of times values are differenced. - label : hashable, optional + label : Hashable, default: "upper" The new coordinate in dimension ``dim`` will have the values of either the minuend's or subtrahend's coordinate for values 'upper' and 'lower', respectively. Other @@ -3227,7 +3582,7 @@ def diff(self, dim: Hashable, n: int = 1, label: Hashable = "upper") -> DataArra Returns ------- - difference : same type as caller + difference : DataArray The n-th order finite difference of this object. Notes @@ -3257,11 +3612,11 @@ def diff(self, dim: Hashable, n: int = 1, label: Hashable = "upper") -> DataArra return self._from_temp_dataset(ds) def shift( - self, - shifts: Mapping[Any, int] = None, + self: T_DataArray, + shifts: Mapping[Any, int] | None = None, fill_value: Any = dtypes.NA, **shifts_kwargs: int, - ) -> DataArray: + ) -> T_DataArray: """Shift this DataArray by an offset along one or more dimensions. Only the data is moved; coordinates stay in place. This is consistent @@ -3273,7 +3628,7 @@ def shift( Parameters ---------- - shifts : mapping of hashable to int, optional + shifts : mapping of Hashable to int or None, optional Integer offset to shift along each of the given dimensions. Positive offsets shift to the right; negative offsets shift to the left. @@ -3307,11 +3662,11 @@ def shift( return self._replace(variable=variable) def roll( - self, - shifts: Mapping[Hashable, int] = None, + self: T_DataArray, + shifts: Mapping[Hashable, int] | None = None, roll_coords: bool = False, **shifts_kwargs: int, - ) -> DataArray: + ) -> T_DataArray: """Roll this array by an offset along one or more dimensions. Unlike shift, roll treats the given dimensions as periodic, so will not @@ -3323,7 +3678,7 @@ def roll( Parameters ---------- - shifts : mapping of hashable to int, optional + shifts : mapping of Hashable to int, optional Integer offset to rotate each of the given dimensions. Positive offsets roll to the right; negative offsets roll to the left. @@ -3356,16 +3711,18 @@ def roll( return self._from_temp_dataset(ds) @property - def real(self) -> DataArray: + def real(self: T_DataArray) -> T_DataArray: return self._replace(self.variable.real) @property - def imag(self) -> DataArray: + def imag(self: T_DataArray) -> T_DataArray: return self._replace(self.variable.imag) def dot( - self, other: DataArray, dims: Hashable | Sequence[Hashable] | None = None - ) -> DataArray: + self: T_DataArray, + other: T_DataArray, + dims: Hashable | Sequence[Hashable] | None = None, + ) -> T_DataArray: """Perform dot product of two DataArrays along their shared dims. Equivalent to taking taking tensordot over all shared dims. @@ -3374,7 +3731,7 @@ def dot( ---------- other : DataArray The other array with which the dot product is performed. - dims : ..., hashable or sequence of hashable, optional + dims : ..., Hashable or sequence of Hashable, optional Which dimensions to sum over. Ellipsis (`...`) sums over all dimensions. If not specified, then all the common dimensions are summed over. @@ -3415,6 +3772,8 @@ def dot( return computation.dot(self, other, dims=dims) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def sortby( self, variables: Hashable | DataArray | Sequence[Hashable | DataArray], @@ -3438,10 +3797,10 @@ def sortby( Parameters ---------- - variables : hashable, DataArray, or sequence of hashable or DataArray + variables : Hashable, DataArray, or sequence of Hashable or DataArray 1D DataArray objects or name(s) of 1D variable(s) in coords whose values are used to sort this array. - ascending : bool, optional + ascending : bool, default: True Whether to sort by ascending or descending order. Returns @@ -3480,14 +3839,14 @@ def sortby( return self._from_temp_dataset(ds) def quantile( - self, + self: T_DataArray, q: ArrayLike, dim: str | Sequence[Hashable] | None = None, method: QUANTILE_METHODS = "linear", - keep_attrs: bool = None, - skipna: bool = None, + keep_attrs: bool | None = None, + skipna: bool | None = None, interpolation: QUANTILE_METHODS = None, - ) -> DataArray: + ) -> T_DataArray: """Compute the qth quantile of the data along the specified dimension. Returns the qth quantiles(s) of the array elements. @@ -3496,7 +3855,7 @@ def quantile( ---------- q : float or array-like of float Quantile to compute, which must be between 0 and 1 inclusive. - dim : hashable or sequence of hashable, optional + dim : Hashable or sequence of Hashable, optional Dimension(s) over which to apply quantile. method : str, default: "linear" This optional parameter specifies the interpolation method to use when the @@ -3526,11 +3885,11 @@ def quantile( previously called "interpolation", renamed in accordance with numpy version 1.22.0. - keep_attrs : bool, optional + keep_attrs : bool or None, optional If True, the dataset's attributes (`attrs`) will be copied from the original object to the new one. If False (default), the new object will be returned without attributes. - skipna : bool, optional + skipna : bool or None, optional If True, skip missing values (as marked by NaN). By default, only skips missing values for float dtypes; other dtypes either do not have a sentinel missing value (int) or skipna=True has not been @@ -3599,8 +3958,11 @@ def quantile( return self._from_temp_dataset(ds) def rank( - self, dim: Hashable, pct: bool = False, keep_attrs: bool = None - ) -> DataArray: + self: T_DataArray, + dim: Hashable, + pct: bool = False, + keep_attrs: bool | None = None, + ) -> T_DataArray: """Ranks the data. Equal values are assigned a rank that is the average of the ranks that @@ -3613,11 +3975,11 @@ def rank( Parameters ---------- - dim : hashable + dim : Hashable Dimension over which to compute rank. - pct : bool, optional + pct : bool, default: False If True, compute percentage ranks, otherwise compute integer ranks. - keep_attrs : bool, optional + keep_attrs : bool or None, optional If True, the dataset's attributes (`attrs`) will be copied from the original object to the new one. If False (default), the new object will be returned without attributes. @@ -3640,8 +4002,11 @@ def rank( return self._from_temp_dataset(ds) def differentiate( - self, coord: Hashable, edge_order: int = 1, datetime_unit: str = None - ) -> DataArray: + self: T_DataArray, + coord: Hashable, + edge_order: Literal[1, 2] = 1, + datetime_unit: DatetimeUnitOptions | None = None, + ) -> T_DataArray: """ Differentiate the array with the second order accurate central differences. @@ -3651,7 +4016,7 @@ def differentiate( Parameters ---------- - coord : hashable + coord : Hashable The coordinate to be used to compute the gradient. edge_order : {1, 2}, default: 1 N-th order accurate differences at the boundaries. @@ -3698,10 +4063,12 @@ def differentiate( ds = self._to_temp_dataset().differentiate(coord, edge_order, datetime_unit) return self._from_temp_dataset(ds) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def integrate( self, coord: Hashable | Sequence[Hashable] = None, - datetime_unit: str = None, + datetime_unit: DatetimeUnitOptions | None = None, ) -> DataArray: """Integrate along the given coordinate using the trapezoidal rule. @@ -3711,7 +4078,7 @@ def integrate( Parameters ---------- - coord : hashable, or sequence of hashable + coord : Hashable, or sequence of Hashable Coordinate(s) used for the integration. datetime_unit : {'Y', 'M', 'W', 'D', 'h', 'm', 's', 'ms', 'us', 'ns', \ 'ps', 'fs', 'as'}, optional @@ -3752,10 +4119,12 @@ def integrate( ds = self._to_temp_dataset().integrate(coord, datetime_unit) return self._from_temp_dataset(ds) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def cumulative_integrate( self, coord: Hashable | Sequence[Hashable] = None, - datetime_unit: str = None, + datetime_unit: DatetimeUnitOptions | None = None, ) -> DataArray: """Integrate cumulatively along the given coordinate using the trapezoidal rule. @@ -3768,7 +4137,7 @@ def cumulative_integrate( Parameters ---------- - coord : hashable, or sequence of hashable + coord : Hashable, or sequence of Hashable Coordinate(s) used for the integration. datetime_unit : {'Y', 'M', 'W', 'D', 'h', 'm', 's', 'ms', 'us', 'ns', \ 'ps', 'fs', 'as'}, optional @@ -3940,8 +4309,8 @@ def polyfit( rcond: float | None = None, w: Hashable | Any | None = None, full: bool = False, - cov: bool = False, - ): + cov: bool | Literal["unscaled"] = False, + ) -> Dataset: """ Least squares polynomial fit. @@ -3950,23 +4319,23 @@ def polyfit( Parameters ---------- - dim : hashable + dim : Hashable Coordinate along which to fit the polynomials. deg : int Degree of the fitting polynomial. - skipna : bool, optional + skipna : bool or None, optional If True, removes all invalid values before fitting each 1D slices of the array. Default is True if data is stored in a dask.array or if there is any invalid values, False otherwise. - rcond : float, optional + rcond : float or None, optional Relative condition number to the fit. - w : hashable or array-like, optional + w : Hashable, array-like or None, optional Weights to apply to the y-coordinate of the sample points. Can be an array-like object or the name of a coordinate in the dataset. - full : bool, optional + full : bool, default: False Whether to return the residuals, matrix rank and singular values in addition to the coefficients. - cov : bool or str, optional + cov : bool or "unscaled", default: False Whether to return to the covariance matrix in addition to the coefficients. The matrix is not scaled if `cov='unscaled'`. @@ -3998,19 +4367,21 @@ def polyfit( ) def pad( - self, + self: T_DataArray, pad_width: Mapping[Any, int | tuple[int, int]] | None = None, - mode: str = "constant", + mode: PadModeOptions = "constant", stat_length: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, - constant_values: (int | tuple[int, int] | Mapping[Any, tuple[int, int]]) + constant_values: float + | tuple[float, float] + | Mapping[Any, tuple[float, float]] | None = None, end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, - reflect_type: str | None = None, + reflect_type: PadReflectOptions = None, **pad_width_kwargs: Any, - ) -> DataArray: + ) -> T_DataArray: """Pad this array along one or more dimensions. .. warning:: @@ -4023,44 +4394,35 @@ def pad( Parameters ---------- - pad_width : mapping of hashable to tuple of int + pad_width : mapping of Hashable to tuple of int Mapping with the form of {dim: (pad_before, pad_after)} describing the number of values padded along each dimension. {dim: pad} is a shortcut for pad_before = pad_after = pad - mode : str, default: "constant" - One of the following string values (taken from numpy docs) - - 'constant' (default) - Pads with a constant value. - 'edge' - Pads with the edge values of array. - 'linear_ramp' - Pads with the linear ramp between end_value and the - array edge value. - 'maximum' - Pads with the maximum value of all or part of the - vector along each axis. - 'mean' - Pads with the mean value of all or part of the - vector along each axis. - 'median' - Pads with the median value of all or part of the - vector along each axis. - 'minimum' - Pads with the minimum value of all or part of the - vector along each axis. - 'reflect' - Pads with the reflection of the vector mirrored on - the first and last values of the vector along each - axis. - 'symmetric' - Pads with the reflection of the vector mirrored - along the edge of the array. - 'wrap' - Pads with the wrap of the vector along the axis. - The first values are used to pad the end and the - end values are used to pad the beginning. - stat_length : int, tuple or mapping of hashable to tuple, default: None + mode : {"constant", "edge", "linear_ramp", "maximum", "mean", "median", \ + "minimum", "reflect", "symmetric", "wrap"}, default: "constant" + How to pad the DataArray (taken from numpy docs): + + - "constant": Pads with a constant value. + - "edge": Pads with the edge values of array. + - "linear_ramp": Pads with the linear ramp between end_value and the + array edge value. + - "maximum": Pads with the maximum value of all or part of the + vector along each axis. + - "mean": Pads with the mean value of all or part of the + vector along each axis. + - "median": Pads with the median value of all or part of the + vector along each axis. + - "minimum": Pads with the minimum value of all or part of the + vector along each axis. + - "reflect": Pads with the reflection of the vector mirrored on + the first and last values of the vector along each axis. + - "symmetric": Pads with the reflection of the vector mirrored + along the edge of the array. + - "wrap": Pads with the wrap of the vector along the axis. + The first values are used to pad the end and the + end values are used to pad the beginning. + + stat_length : int, tuple or mapping of Hashable to tuple, default: None Used in 'maximum', 'mean', 'median', and 'minimum'. Number of values at edge of each axis used to calculate the statistic value. {dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)} unique @@ -4070,7 +4432,7 @@ def pad( (stat_length,) or int is a shortcut for before = after = statistic length for all axes. Default is ``None``, to use the entire axis. - constant_values : scalar, tuple or mapping of hashable to tuple, default: 0 + constant_values : scalar, tuple or mapping of Hashable to tuple, default: 0 Used in 'constant'. The values to set the padded values for each axis. ``{dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)}`` unique @@ -4080,7 +4442,7 @@ def pad( ``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for all dimensions. Default is 0. - end_values : scalar, tuple or mapping of hashable to tuple, default: 0 + end_values : scalar, tuple or mapping of Hashable to tuple, default: 0 Used in 'linear_ramp'. The values used for the ending value of the linear_ramp and that will form the edge of the padded array. ``{dim_1: (before_1, after_1), ... dim_N: (before_N, after_N)}`` unique @@ -4090,9 +4452,9 @@ def pad( ``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for all axes. Default is 0. - reflect_type : {"even", "odd"}, optional - Used in "reflect", and "symmetric". The "even" style is the - default with an unaltered reflection around the edge value. For + reflect_type : {"even", "odd", None}, optional + Used in "reflect", and "symmetric". The "even" style is the + default with an unaltered reflection around the edge value. For the "odd" style, the extended part of the array is created by subtracting the reflected values from two times the edge value. **pad_width_kwargs @@ -4168,10 +4530,10 @@ def pad( def idxmin( self, - dim: Hashable = None, - skipna: bool = None, + dim: Hashable | None = None, + skipna: bool | None = None, fill_value: Any = dtypes.NA, - keep_attrs: bool = None, + keep_attrs: bool | None = None, ) -> DataArray: """Return the coordinate label of the minimum value along a dimension. @@ -4198,9 +4560,9 @@ def idxmin( null. By default this is NaN. The fill value and result are automatically converted to a compatible dtype if possible. Ignored if ``skipna`` is False. - keep_attrs : bool, default: False + keep_attrs : bool or None, optional If True, the attributes (``attrs``) will be copied from the - original object to the new one. If False (default), the new object + original object to the new one. If False, the new object will be returned without attributes. Returns @@ -4265,9 +4627,9 @@ def idxmin( def idxmax( self, dim: Hashable = None, - skipna: bool = None, + skipna: bool | None = None, fill_value: Any = dtypes.NA, - keep_attrs: bool = None, + keep_attrs: bool | None = None, ) -> DataArray: """Return the coordinate label of the maximum value along a dimension. @@ -4280,7 +4642,7 @@ def idxmax( Parameters ---------- - dim : hashable, optional + dim : Hashable, optional Dimension over which to apply `idxmax`. This is optional for 1D arrays, but required for arrays with 2 or more dimensions. skipna : bool or None, default: None @@ -4294,9 +4656,9 @@ def idxmax( null. By default this is NaN. The fill value and result are automatically converted to a compatible dtype if possible. Ignored if ``skipna`` is False. - keep_attrs : bool, default: False + keep_attrs : bool or None, optional If True, the attributes (``attrs``) will be copied from the - original object to the new one. If False (default), the new object + original object to the new one. If False, the new object will be returned without attributes. Returns @@ -4358,12 +4720,14 @@ def idxmax( keep_attrs=keep_attrs, ) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def argmin( self, - dim: Hashable | Sequence[Hashable] = None, - axis: int = None, - keep_attrs: bool = None, - skipna: bool = None, + dim: Hashable | Sequence[Hashable] | None = None, + axis: int | None = None, + keep_attrs: bool | None = None, + skipna: bool | None = None, ) -> DataArray | dict[Hashable, DataArray]: """Index or indices of the minimum of the DataArray over one or more dimensions. @@ -4376,19 +4740,19 @@ def argmin( Parameters ---------- - dim : hashable, sequence of hashable or ..., optional + dim : Hashable, sequence of Hashable, None or ..., optional The dimensions over which to find the minimum. By default, finds minimum over all dimensions - for now returning an int for backward compatibility, but this is deprecated, in future will return a dict with indices for all dimensions; to return a dict with all dimensions now, pass '...'. - axis : int, optional + axis : int or None, optional Axis over which to apply `argmin`. Only one of the 'dim' and 'axis' arguments can be supplied. - keep_attrs : bool, optional + keep_attrs : bool or None, optional If True, the attributes (`attrs`) will be copied from the original - object to the new one. If False (default), the new object will be + object to the new one. If False, the new object will be returned without attributes. - skipna : bool, optional + skipna : bool or None, optional If True, skip missing values (as marked by NaN). By default, only skips missing values for float dtypes; other dtypes either do not have a sentinel missing value (int) or skipna=True has not been @@ -4461,12 +4825,14 @@ def argmin( else: return self._replace_maybe_drop_dims(result) + # change type of self and return to T_DataArray once + # https://github.com/python/mypy/issues/12846 is resolved def argmax( self, dim: Hashable | Sequence[Hashable] = None, - axis: int = None, - keep_attrs: bool = None, - skipna: bool = None, + axis: int | None = None, + keep_attrs: bool | None = None, + skipna: bool | None = None, ) -> DataArray | dict[Hashable, DataArray]: """Index or indices of the maximum of the DataArray over one or more dimensions. @@ -4479,19 +4845,19 @@ def argmax( Parameters ---------- - dim : hashable, sequence of hashable or ..., optional + dim : Hashable, sequence of Hashable, None or ..., optional The dimensions over which to find the maximum. By default, finds maximum over all dimensions - for now returning an int for backward compatibility, but this is deprecated, in future will return a dict with indices for all dimensions; to return a dict with all dimensions now, pass '...'. - axis : int, optional + axis : int or None, optional Axis over which to apply `argmax`. Only one of the 'dim' and 'axis' arguments can be supplied. - keep_attrs : bool, optional + keep_attrs : bool or None, optional If True, the attributes (`attrs`) will be copied from the original - object to the new one. If False (default), the new object will be + object to the new one. If False, the new object will be returned without attributes. - skipna : bool, optional + skipna : bool or None, optional If True, skip missing values (as marked by NaN). By default, only skips missing values for float dtypes; other dtypes either do not have a sentinel missing value (int) or skipna=True has not been @@ -4566,10 +4932,10 @@ def argmax( def query( self, - queries: Mapping[Any, Any] = None, - parser: str = "pandas", - engine: str = None, - missing_dims: str = "raise", + queries: Mapping[Any, Any] | None = None, + parser: QueryParserOptions = "pandas", + engine: QueryEngineOptions = None, + missing_dims: ErrorOptionsWithWarn = "raise", **queries_kwargs: Any, ) -> DataArray: """Return a new data array indexed along the specified @@ -4578,8 +4944,8 @@ def query( Parameters ---------- - queries : dict, optional - A dict with keys matching dimensions and values given by strings + queries : dict-like or None, optional + A dict-like with keys matching dimensions and values given by strings containing Python expressions to be evaluated against the data variables in the dataset. The expressions will be evaluated using the pandas eval() function, and can contain any valid Python expressions but cannot @@ -4591,15 +4957,19 @@ def query( parser to retain strict Python semantics. engine : {"python", "numexpr", None}, default: None The engine used to evaluate the expression. Supported engines are: + - None: tries to use numexpr, falls back to python - "numexpr": evaluates expressions using numexpr - "python": performs operations as if you had eval’d in top level python + missing_dims : {"raise", "warn", "ignore"}, default: "raise" What to do if dimensions that should be selected from are not present in the - Dataset: + DataArray: + - "raise": raise an exception - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions + **queries_kwargs : {dim: query, ...}, optional The keyword arguments form of ``queries``. One of queries or queries_kwargs must be provided. @@ -4643,13 +5013,13 @@ def curvefit( self, coords: str | DataArray | Iterable[str | DataArray], func: Callable[..., Any], - reduce_dims: Hashable | Iterable[Hashable] = None, + reduce_dims: Hashable | Iterable[Hashable] | None = None, skipna: bool = True, - p0: dict[str, Any] = None, - bounds: dict[str, Any] = None, - param_names: Sequence[str] = None, - kwargs: dict[str, Any] = None, - ): + p0: dict[str, Any] | None = None, + bounds: dict[str, Any] | None = None, + param_names: Sequence[str] | None = None, + kwargs: dict[str, Any] | None = None, + ) -> Dataset: """ Curve fitting optimization for arbitrary functions. @@ -4657,7 +5027,7 @@ def curvefit( Parameters ---------- - coords : hashable, DataArray, or sequence of DataArray or hashable + coords : Hashable, DataArray, or sequence of DataArray or Hashable Independent coordinate(s) over which to perform the curve fitting. Must share at least one dimension with the calling object. When fitting multi-dimensional functions, supply `coords` as a sequence in the same order as arguments in @@ -4668,22 +5038,22 @@ def curvefit( array of length `len(x)`. `params` are the fittable parameters which are optimized by scipy curve_fit. `x` can also be specified as a sequence containing multiple coordinates, e.g. `f((x0, x1), *params)`. - reduce_dims : hashable or sequence of hashable + reduce_dims : Hashable or sequence of Hashable Additional dimension(s) over which to aggregate while fitting. For example, calling `ds.curvefit(coords='time', reduce_dims=['lat', 'lon'], ...)` will aggregate all lat and lon points and fit the specified function along the time dimension. - skipna : bool, optional + skipna : bool, default: True Whether to skip missing values when fitting. Default is True. - p0 : dict-like, optional + p0 : dict-like or None, optional Optional dictionary of parameter names to initial guesses passed to the `curve_fit` `p0` arg. If none or only some parameters are passed, the rest will be assigned initial values following the default scipy behavior. - bounds : dict-like, optional + bounds : dict-like or None, optional Optional dictionary of parameter names to bounding values passed to the `curve_fit` `bounds` arg. If none or only some parameters are passed, the rest will be unbounded following the default scipy behavior. - param_names : sequence of hashable, optional + param_names : sequence of Hashable or None, optional Sequence of names for the fittable parameters of `func`. If not supplied, this will be automatically determined by arguments of `func`. `param_names` should be manually supplied when fitting a function that takes a variable @@ -4718,10 +5088,10 @@ def curvefit( ) def drop_duplicates( - self, - dim: Hashable | Iterable[Hashable] | ..., - keep: Literal["first", "last"] | Literal[False] = "first", - ): + self: T_DataArray, + dim: Hashable | Iterable[Hashable], + keep: Literal["first", "last", False] = "first", + ) -> T_DataArray: """Returns a new DataArray with duplicate dimension values removed. Parameters @@ -4730,6 +5100,7 @@ def drop_duplicates( Pass `...` to drop duplicates along all dimensions. keep : {"first", "last", False}, default: "first" Determines which duplicates (if any) to keep. + - ``"first"`` : Drop duplicates except for the first occurrence. - ``"last"`` : Drop duplicates except for the last occurrence. - False : Drop all duplicates. @@ -4901,4 +5272,4 @@ def interp_calendar( # this needs to be at the end, or mypy will confuse with `str` # https://mypy.readthedocs.io/en/latest/common_issues.html#dealing-with-conflicting-names - str = utils.UncachedAccessor(StringAccessor) + str = utils.UncachedAccessor(StringAccessor["DataArray"]) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 855718cfe74..38e4d2eadef 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3,6 +3,7 @@ import copy import datetime import inspect +import itertools import sys import warnings from collections import defaultdict @@ -29,8 +30,6 @@ import numpy as np import pandas as pd -import xarray as xr - from ..coding.calendar_ops import convert_calendar, interp_calendar from ..coding.cftimeindex import CFTimeIndex, _parse_array_of_cftime_strings from ..plot.dataset_plot import _Dataset_PlotMethods @@ -102,14 +101,28 @@ if TYPE_CHECKING: from ..backends import AbstractDataStore, ZarrStore + from ..backends.api import T_NetcdfEngine, T_NetcdfTypes from .dataarray import DataArray from .merge import CoercibleMapping - from .types import T_Xarray + from .types import ( + CombineAttrsOptions, + CompatOptions, + ErrorOptions, + ErrorOptionsWithWarn, + InterpOptions, + JoinOptions, + PadModeOptions, + PadReflectOptions, + QueryEngineOptions, + QueryParserOptions, + T_Dataset, + T_Xarray, + ) try: from dask.delayed import Delayed except ImportError: - Delayed = None + Delayed = None # type: ignore # list of attributes of pd.DatetimeIndex that are ndarrays of time info @@ -138,6 +151,8 @@ def _get_virtual_variable( objects (if possible) """ + from .dataarray import DataArray + if dim_sizes is None: dim_sizes = {} @@ -157,7 +172,7 @@ def _get_virtual_variable( ref_var = variables[ref_name] if _contains_datetime_like_objects(ref_var): - ref_var = xr.DataArray(ref_var) + ref_var = DataArray(ref_var) data = getattr(ref_var.dt, var_name).data else: data = getattr(ref_var, var_name).data @@ -171,60 +186,63 @@ def _assert_empty(args: tuple, msg: str = "%s") -> None: raise ValueError(msg % args) -def _check_chunks_compatibility(var, chunks, preferred_chunks): - for dim in var.dims: - if dim not in chunks or (dim not in preferred_chunks): - continue - - preferred_chunks_dim = preferred_chunks.get(dim) - chunks_dim = chunks.get(dim) - - if isinstance(chunks_dim, int): - chunks_dim = (chunks_dim,) - else: - chunks_dim = chunks_dim[:-1] - - if any(s % preferred_chunks_dim for s in chunks_dim): - warnings.warn( - f"Specified Dask chunks {chunks[dim]} would separate " - f"on disks chunk shape {preferred_chunks[dim]} for dimension {dim}. " - "This could degrade performance. " - "Consider rechunking after loading instead.", - stacklevel=2, - ) - - def _get_chunk(var, chunks): - # chunks need to be explicitly computed to take correctly into account - # backend preferred chunking + """ + Return map from each dim to chunk sizes, accounting for backend's preferred chunks. + """ + import dask.array as da if isinstance(var, IndexVariable): return {} + dims = var.dims + shape = var.shape - if isinstance(chunks, int) or (chunks == "auto"): - chunks = dict.fromkeys(var.dims, chunks) - + # Determine the explicit requested chunks. preferred_chunks = var.encoding.get("preferred_chunks", {}) - preferred_chunks_list = [ - preferred_chunks.get(dim, shape) for dim, shape in zip(var.dims, var.shape) - ] - - chunks_list = [ - chunks.get(dim, None) or preferred_chunks.get(dim, None) for dim in var.dims - ] - - output_chunks_list = da.core.normalize_chunks( - chunks_list, - shape=var.shape, - dtype=var.dtype, - previous_chunks=preferred_chunks_list, + preferred_chunk_shape = tuple( + preferred_chunks.get(dim, size) for dim, size in zip(dims, shape) + ) + if isinstance(chunks, Number) or (chunks == "auto"): + chunks = dict.fromkeys(dims, chunks) + chunk_shape = tuple( + chunks.get(dim, None) or preferred_chunk_sizes + for dim, preferred_chunk_sizes in zip(dims, preferred_chunk_shape) + ) + chunk_shape = da.core.normalize_chunks( + chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape ) - output_chunks = dict(zip(var.dims, output_chunks_list)) - _check_chunks_compatibility(var, output_chunks, preferred_chunks) + # Warn where requested chunks break preferred chunks, provided that the variable + # contains data. + if var.size: + for dim, size, chunk_sizes in zip(dims, shape, chunk_shape): + try: + preferred_chunk_sizes = preferred_chunks[dim] + except KeyError: + continue + # Determine the stop indices of the preferred chunks, but omit the last stop + # (equal to the dim size). In particular, assume that when a sequence + # expresses the preferred chunks, the sequence sums to the size. + preferred_stops = ( + range(preferred_chunk_sizes, size, preferred_chunk_sizes) + if isinstance(preferred_chunk_sizes, Number) + else itertools.accumulate(preferred_chunk_sizes[:-1]) + ) + # Gather any stop indices of the specified chunks that are not a stop index + # of a preferred chunk. Again, omit the last stop, assuming that it equals + # the dim size. + breaks = set(itertools.accumulate(chunk_sizes[:-1])).difference( + preferred_stops + ) + if breaks: + warnings.warn( + "The specified Dask chunks separate the stored chunks along " + f'dimension "{dim}" starting at index {min(breaks)}. This could ' + "degrade performance. Instead, consider rechunking after loading." + ) - return output_chunks + return dict(zip(dims, chunk_shape)) def _maybe_chunk( @@ -235,6 +253,7 @@ def _maybe_chunk( lock=None, name_prefix="xarray-", overwrite_encoded_chunks=False, + inline_array=False, ): from dask.base import tokenize @@ -246,7 +265,7 @@ def _maybe_chunk( # subtle bugs result otherwise. see GH3350 token2 = tokenize(name, token if token else var._data, chunks) name2 = f"{name_prefix}{name}-{token2}" - var = var.chunk(chunks, name=name2, lock=lock) + var = var.chunk(chunks, name=name2, lock=lock, inline_array=inline_array) if overwrite_encoded_chunks and var.chunks is not None: var.encoding["chunks"] = tuple(x[0] for x in var.chunks) @@ -1366,7 +1385,8 @@ def __setitem__(self, key: Hashable | list[Hashable] | Mapping, value) -> None: If the given value is also a dataset, select corresponding variables in the given value and in the dataset to be changed. - If value is a `DataArray`, call its `select_vars()` method, rename it + If value is a ` + from .dataarray import DataArray`, call its `select_vars()` method, rename it to `key` and merge the contents of the resulting dataset into this dataset. @@ -1374,6 +1394,8 @@ def __setitem__(self, key: Hashable | list[Hashable] | Mapping, value) -> None: ``(dims, data[, attrs])``), add it to this dataset as a new variable. """ + from .dataarray import DataArray + if utils.is_dict_like(key): # check for consistency and convert value to dataset value = self._setitem_check(key, value) @@ -1407,7 +1429,7 @@ def __setitem__(self, key: Hashable | list[Hashable] | Mapping, value) -> None: ) if isinstance(value, Dataset): self.update(dict(zip(key, value.data_vars.values()))) - elif isinstance(value, xr.DataArray): + elif isinstance(value, DataArray): raise ValueError("Cannot assign single DataArray to multiple keys") else: self.update(dict(zip(key, value))) @@ -1426,6 +1448,7 @@ def _setitem_check(self, key, value): When assigning values to a subset of a Dataset, do consistency check beforehand to avoid leaving the dataset in a partially updated state when an error occurs. """ + from .alignment import align from .dataarray import DataArray if isinstance(value, Dataset): @@ -1442,7 +1465,7 @@ def _setitem_check(self, key, value): "Dataset assignment only accepts DataArrays, Datasets, and scalars." ) - new_value = xr.Dataset() + new_value = Dataset() for name, var in self.items(): # test indexing try: @@ -1479,7 +1502,7 @@ def _setitem_check(self, key, value): # check consistency of dimension sizes and dimension coordinates if isinstance(value, DataArray) or isinstance(value, Dataset): - xr.align(self[key], value, join="exact", copy=False) + align(self[key], value, join="exact", copy=False) return new_value @@ -1671,15 +1694,64 @@ def dump_to_store(self, store: AbstractDataStore, **kwargs) -> None: # with to_netcdf() dump_to_store(self, store, **kwargs) + # path=None writes to bytes + @overload + def to_netcdf( + self, + path: None = None, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: bool = True, + invalid_netcdf: bool = False, + ) -> bytes: + ... + + # default return None + @overload + def to_netcdf( + self, + path: str | PathLike, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + compute: Literal[True] = True, + invalid_netcdf: bool = False, + ) -> None: + ... + + # compute=False returns dask.Delayed + @overload + def to_netcdf( + self, + path: str | PathLike, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, + *, + compute: Literal[False], + invalid_netcdf: bool = False, + ) -> Delayed: + ... + def to_netcdf( self, - path=None, - mode: str = "w", - format: str = None, - group: str = None, - engine: str = None, - encoding: Mapping = None, - unlimited_dims: Iterable[Hashable] = None, + path: str | PathLike | None = None, + mode: Literal["w", "a"] = "w", + format: T_NetcdfTypes | None = None, + group: str | None = None, + engine: T_NetcdfEngine | None = None, + encoding: Mapping[Hashable, Mapping[str, Any]] | None = None, + unlimited_dims: Iterable[Hashable] | None = None, compute: bool = True, invalid_netcdf: bool = False, ) -> bytes | Delayed | None: @@ -1749,39 +1821,89 @@ def to_netcdf( Only valid along with ``engine="h5netcdf"``. If True, allow writing hdf5 files which are invalid netcdf as described in https://github.com/h5netcdf/h5netcdf. + + Returns + ------- + * ``bytes`` if path is None + * ``dask.delayed.Delayed`` if compute is False + * None otherwise + + See Also + -------- + DataArray.to_netcdf """ if encoding is None: encoding = {} from ..backends.api import to_netcdf - return to_netcdf( + return to_netcdf( # type: ignore # mypy cannot resolve the overloads:( self, path, - mode, + mode=mode, format=format, group=group, engine=engine, encoding=encoding, unlimited_dims=unlimited_dims, compute=compute, + multifile=False, invalid_netcdf=invalid_netcdf, ) + # compute=True (default) returns ZarrStore + @overload def to_zarr( self, - store: MutableMapping | str | PathLike | None = None, + store: MutableMapping | str | PathLike[str] | None = None, chunk_store: MutableMapping | str | PathLike | None = None, - mode: str = None, + mode: Literal["w", "w-", "a", "r+", None] = None, synchronizer=None, - group: str = None, - encoding: Mapping = None, - compute: bool = True, + group: str | None = None, + encoding: Mapping | None = None, + compute: Literal[True] = True, consolidated: bool | None = None, - append_dim: Hashable = None, - region: Mapping[str, slice] = None, + append_dim: Hashable | None = None, + region: Mapping[str, slice] | None = None, safe_chunks: bool = True, - storage_options: dict[str, str] = None, + storage_options: dict[str, str] | None = None, ) -> ZarrStore: + ... + + # compute=False returns dask.Delayed + @overload + def to_zarr( + self, + store: MutableMapping | str | PathLike[str] | None = None, + chunk_store: MutableMapping | str | PathLike | None = None, + mode: Literal["w", "w-", "a", "r+", None] = None, + synchronizer=None, + group: str | None = None, + encoding: Mapping | None = None, + *, + compute: Literal[False], + consolidated: bool | None = None, + append_dim: Hashable | None = None, + region: Mapping[str, slice] | None = None, + safe_chunks: bool = True, + storage_options: dict[str, str] | None = None, + ) -> Delayed: + ... + + def to_zarr( + self, + store: MutableMapping | str | PathLike[str] | None = None, + chunk_store: MutableMapping | str | PathLike | None = None, + mode: Literal["w", "w-", "a", "r+", None] = None, + synchronizer=None, + group: str | None = None, + encoding: Mapping | None = None, + compute: bool = True, + consolidated: bool | None = None, + append_dim: Hashable | None = None, + region: Mapping[str, slice] | None = None, + safe_chunks: bool = True, + storage_options: dict[str, str] | None = None, + ) -> ZarrStore | Delayed: """Write dataset contents to a zarr group. Zarr chunks are determined in the following way: @@ -1862,6 +1984,11 @@ def to_zarr( Any additional parameters for the storage backend (ignored for local paths). + Returns + ------- + * ``dask.delayed.Delayed`` if compute is False + * ZarrStore otherwise + References ---------- https://zarr.readthedocs.io/ @@ -1886,10 +2013,7 @@ def to_zarr( """ from ..backends.api import to_zarr - if encoding is None: - encoding = {} - - return to_zarr( + return to_zarr( # type: ignore self, store=store, chunk_store=chunk_store, @@ -1990,6 +2114,8 @@ def chunk( name_prefix: str = "xarray-", token: str = None, lock: bool = False, + inline_array: bool = False, + **chunks_kwargs: Any, ) -> Dataset: """Coerce all arrays in this dataset into dask arrays with the given chunks. @@ -2003,7 +2129,7 @@ def chunk( Parameters ---------- - chunks : int, "auto" or mapping of hashable to int, optional + chunks : int, tuple of int, "auto" or mapping of hashable to int, optional Chunk sizes along each dimension, e.g., ``5``, ``"auto"``, or ``{"x": 5, "y": 5}``. name_prefix : str, optional @@ -2013,6 +2139,12 @@ def chunk( lock : optional Passed on to :py:func:`dask.array.from_array`, if the array is not already as dask array. + inline_array: optional + Passed on to :py:func:`dask.array.from_array`, if the array is not + already as dask array. + **chunks_kwargs : {dim: chunks, ...}, optional + The keyword arguments form of ``chunks``. + One of chunks or chunks_kwargs must be provided Returns ------- @@ -2023,8 +2155,9 @@ def chunk( Dataset.chunks Dataset.chunksizes xarray.unify_chunks + dask.array.from_array """ - if chunks is None: + if chunks is None and chunks_kwargs is None: warnings.warn( "None value for 'chunks' is deprecated. " "It will raise an error in the future. Use instead '{}'", @@ -2034,6 +2167,8 @@ def chunk( if isinstance(chunks, (Number, str, int)): chunks = dict.fromkeys(self.dims, chunks) + else: + chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk") bad_dims = chunks.keys() - self.dims.keys() if bad_dims: @@ -2048,7 +2183,7 @@ def chunk( return self._replace(variables) def _validate_indexers( - self, indexers: Mapping[Any, Any], missing_dims: str = "raise" + self, indexers: Mapping[Any, Any], missing_dims: ErrorOptionsWithWarn = "raise" ) -> Iterator[tuple[Hashable, int | slice | np.ndarray | Variable]]: """Here we make sure + indexer has a valid keys @@ -2056,6 +2191,7 @@ def _validate_indexers( + string indexers are cast to the appropriate date type if the associated index is a DatetimeIndex or CFTimeIndex """ + from ..coding.cftimeindex import CFTimeIndex from .dataarray import DataArray indexers = drop_dims_from_indexers(indexers, self.dims, missing_dims) @@ -2079,7 +2215,7 @@ def _validate_indexers( index = self._indexes[k].to_pandas_index() if isinstance(index, pd.DatetimeIndex): v = v.astype("datetime64[ns]") - elif isinstance(index, xr.CFTimeIndex): + elif isinstance(index, CFTimeIndex): v = _parse_array_of_cftime_strings(v, index.date_type) if v.ndim > 1: @@ -2153,7 +2289,7 @@ def isel( self, indexers: Mapping[Any, Any] = None, drop: bool = False, - missing_dims: str = "raise", + missing_dims: ErrorOptionsWithWarn = "raise", **indexers_kwargs: Any, ) -> Dataset: """Returns a new dataset with each array indexed along the specified @@ -2172,14 +2308,14 @@ def isel( If DataArrays are passed as indexers, xarray-style indexing will be carried out. See :ref:`indexing` for the details. One of indexers or indexers_kwargs must be provided. - drop : bool, optional + drop : bool, default: False If ``drop=True``, drop coordinates variables indexed by integers instead of making them scalar. missing_dims : {"raise", "warn", "ignore"}, default: "raise" What to do if dimensions that should be selected from are not present in the Dataset: - "raise": raise an exception - - "warning": raise a warning, and ignore the missing dimensions + - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions **indexers_kwargs : {dim: indexer, ...}, optional The keyword arguments form of ``indexers``. @@ -2244,7 +2380,7 @@ def _isel_fancy( indexers: Mapping[Any, Any], *, drop: bool, - missing_dims: str = "raise", + missing_dims: ErrorOptionsWithWarn = "raise", ) -> Dataset: valid_indexers = dict(self._validate_indexers(indexers, missing_dims)) @@ -2260,6 +2396,10 @@ def _isel_fancy( } if var_indexers: new_var = var.isel(indexers=var_indexers) + # drop scalar coordinates + # https://github.com/pydata/xarray/issues/6554 + if name in self.coords and drop and new_var.ndim == 0: + continue else: new_var = var.copy(deep=False) if name not in indexes: @@ -2507,8 +2647,8 @@ def thin( return self.isel(indexers_slices) def broadcast_like( - self, other: Dataset | DataArray, exclude: Iterable[Hashable] = None - ) -> Dataset: + self: T_Dataset, other: Dataset | DataArray, exclude: Iterable[Hashable] = None + ) -> T_Dataset: """Broadcast this DataArray against another Dataset or DataArray. This is equivalent to xr.broadcast(other, self)[1] @@ -2528,7 +2668,9 @@ def broadcast_like( dims_map, common_coords = _get_broadcast_dims_map_common_coords(args, exclude) - return _broadcast_helper(args[1], exclude, dims_map, common_coords) + return _broadcast_helper( + cast("T_Dataset", args[1]), exclude, dims_map, common_coords + ) def _reindex_callback( self, @@ -2896,13 +3038,22 @@ def _reindex( def interp( self, coords: Mapping[Any, Any] = None, - method: str = "linear", + method: InterpOptions = "linear", assume_sorted: bool = False, kwargs: Mapping[str, Any] = None, method_non_numeric: str = "nearest", **coords_kwargs: Any, ) -> Dataset: - """Multidimensional interpolation of Dataset. + """Interpolate a Dataset onto new coordinates + + Performs univariate or multivariate interpolation of a Dataset onto + new coordinates using scipy's interpolation routines. If interpolating + along an existing dimension, :py:class:`scipy.interpolate.interp1d` is + called. When interpolating along multiple existing dimensions, an + attempt is made to decompose the interpolation into multiple + 1-dimensional interpolations. If this is possible, + :py:class:`scipy.interpolate.interp1d` is called. Otherwise, + :py:func:`scipy.interpolate.interpn` is called. Parameters ---------- @@ -2912,9 +3063,15 @@ def interp( If DataArrays are passed as new coordinates, their dimensions are used for the broadcasting. Missing values are skipped. method : str, optional - {"linear", "nearest"} for multidimensional array, - {"linear", "nearest", "zero", "slinear", "quadratic", "cubic"} - for 1-dimensional array. "linear" is used by default. + The method used to interpolate. The method should be supported by + the scipy interpolator: + + - ``interp1d``: {"linear", "nearest", "zero", "slinear", + "quadratic", "cubic", "polynomial"} + - ``interpn``: {"linear", "nearest"} + + If ``"polynomial"`` is passed, the ``order`` keyword argument must + also be provided. assume_sorted : bool, optional If False, values of coordinates that are interpolated over can be in any order and they are sorted first. If True, interpolated @@ -2922,8 +3079,8 @@ def interp( values. kwargs : dict, optional Additional keyword arguments passed to scipy's interpolator. Valid - options and their behavior depend on if 1-dimensional or - multi-dimensional interpolation is used. + options and their behavior depend whether ``interp1d`` or + ``interpn`` is used. method_non_numeric : {"nearest", "pad", "ffill", "backfill", "bfill"}, optional Method for non-numeric types. Passed on to :py:meth:`Dataset.reindex`. ``"nearest"`` is used by default. @@ -3157,7 +3314,7 @@ def _validate_interp_indexer(x, new_x): def interp_like( self, other: Dataset | DataArray, - method: str = "linear", + method: InterpOptions = "linear", assume_sorted: bool = False, kwargs: Mapping[str, Any] = None, method_non_numeric: str = "nearest", @@ -3165,6 +3322,13 @@ def interp_like( """Interpolate this object onto the coordinates of another object, filling the out of range values with NaN. + If interpolating along a single existing dimension, + :py:class:`scipy.interpolate.interp1d` is called. When interpolating + along multiple existing dimensions, an attempt is made to decompose the + interpolation into multiple 1-dimensional interpolations. If this is + possible, :py:class:`scipy.interpolate.interp1d` is called. Otherwise, + :py:func:`scipy.interpolate.interpn` is called. + Parameters ---------- other : Dataset or DataArray @@ -3172,9 +3336,15 @@ def interp_like( names to an 1d array-like, which provides coordinates upon which to index the variables in this dataset. Missing values are skipped. method : str, optional - {"linear", "nearest"} for multidimensional array, - {"linear", "nearest", "zero", "slinear", "quadratic", "cubic"} - for 1-dimensional array. 'linear' is used by default. + The method used to interpolate. The method should be supported by + the scipy interpolator: + + - {"linear", "nearest", "zero", "slinear", "quadratic", "cubic", + "polynomial"} when ``interp1d`` is called. + - {"linear", "nearest"} when ``interpn`` is called. + + If ``"polynomial"`` is passed, the ``order`` keyword argument must + also be provided. assume_sorted : bool, optional If False, values of coordinates that are interpolated over can be in any order and they are sorted first. If True, interpolated @@ -3533,9 +3703,9 @@ def expand_dims( and the values are either integers (giving the length of the new dimensions) or array-like (giving the coordinates of the new dimensions). - axis : int, sequence of int, or None + axis : int, sequence of int, or None, default: None Axis position(s) where new axis is to be inserted (position(s) on - the result array). If a list (or tuple) of integers is passed, + the result array). If a sequence of integers is passed, multiple axes are inserted. In this case, dim arguments should be same length list. If axis=None is passed, all the axes will be inserted to the start of the result array. @@ -3547,8 +3717,8 @@ def expand_dims( Returns ------- - expanded : same type as caller - This object, but with an additional dimension(s). + expanded : Dataset + This object, but with additional dimension(s). """ if dim is None: pass @@ -3864,18 +4034,18 @@ def reset_index( def reorder_levels( self, - dim_order: Mapping[Any, Sequence[int]] = None, - **dim_order_kwargs: Sequence[int], + dim_order: Mapping[Any, Sequence[int | Hashable]] = None, + **dim_order_kwargs: Sequence[int | Hashable], ) -> Dataset: """Rearrange index levels using input order. Parameters ---------- - dim_order : optional + dim_order : dict-like of Hashable to Sequence of int or Hashable, optional Mapping from names matching dimensions and values given by lists representing new level orders. Every given dimension must have a multi-index. - **dim_order_kwargs : optional + **dim_order_kwargs : Sequence of int or Hashable, optional The keyword arguments form of ``dim_order``. One of dim_order or dim_order_kwargs must be provided. @@ -4040,8 +4210,8 @@ def stack( ellipsis (`...`) will be replaced by all unlisted dimensions. Passing a list containing an ellipsis (`stacked_dim=[...]`) will stack over all dimensions. - create_index : bool, optional - If True (default), create a multi-index for each of the stacked dimensions. + create_index : bool or None, default: True + If True, create a multi-index for each of the stacked dimensions. If False, don't create any index. If None, create a multi-index only if exactly one single (1-d) coordinate index is found for every dimension to stack. @@ -4141,6 +4311,8 @@ def to_stacked_array( Dimensions without coordinates: x """ + from .concat import concat + stacking_dims = tuple(dim for dim in self.dims if dim not in sample_dims) for variable in self: @@ -4171,7 +4343,7 @@ def ensure_stackable(val): # concatenate the arrays stackable_vars = [ensure_stackable(self[key]) for key in self.data_vars] - data_array = xr.concat(stackable_vars, dim=new_dim) + data_array = concat(stackable_vars, dim=new_dim) if name is not None: data_array.name = name @@ -4415,10 +4587,10 @@ def merge( self, other: CoercibleMapping | DataArray, overwrite_vars: Hashable | Iterable[Hashable] = frozenset(), - compat: str = "no_conflicts", - join: str = "outer", + compat: CompatOptions = "no_conflicts", + join: JoinOptions = "outer", fill_value: Any = dtypes.NA, - combine_attrs: str = "override", + combine_attrs: CombineAttrsOptions = "override", ) -> Dataset: """Merge the arrays of two datasets into a single dataset. @@ -4447,6 +4619,7 @@ def merge( - 'no_conflicts': only values which are not null in both datasets must be equal. The returned dataset then contains the combination of all non-null values. + join : {"outer", "inner", "left", "right", "exact"}, optional Method for joining ``self`` and ``other`` along shared dimensions: @@ -4455,12 +4628,14 @@ def merge( - 'left': use indexes from ``self`` - 'right': use indexes from ``other`` - 'exact': error instead of aligning non-equal indexes + fill_value : scalar or dict-like, optional Value to use for newly missing values. If a dict-like, maps variable names (including coordinates) to fill values. combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \ - "override"}, default: "override" - String indicating how to combine attrs of the objects being merged: + "override"} or callable, default: "override" + A callable or a string indicating how to combine attrs of the objects being + merged: - "drop": empty attrs on returned Dataset. - "identical": all attrs must be the same on every object. @@ -4471,6 +4646,9 @@ def merge( - "override": skip comparing and copy attrs from the first dataset to the result. + If a callable, it must expect a sequence of ``attrs`` dicts and a context object + as its only parameters. + Returns ------- merged : Dataset @@ -4485,7 +4663,9 @@ def merge( -------- Dataset.update """ - other = other.to_dataset() if isinstance(other, xr.DataArray) else other + from .dataarray import DataArray + + other = other.to_dataset() if isinstance(other, DataArray) else other merge_result = dataset_merge_method( self, other, @@ -4510,7 +4690,7 @@ def _assert_all_in_dataset( ) def drop_vars( - self, names: Hashable | Iterable[Hashable], *, errors: str = "raise" + self, names: Hashable | Iterable[Hashable], *, errors: ErrorOptions = "raise" ) -> Dataset: """Drop variables from this dataset. @@ -4518,8 +4698,8 @@ def drop_vars( ---------- names : hashable or iterable of hashable Name(s) of variables to drop. - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if any of the variable + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the variable passed are not in the dataset. If 'ignore', any given names that are in the dataset are dropped and no error is raised. @@ -4536,6 +4716,23 @@ def drop_vars( if errors == "raise": self._assert_all_in_dataset(names) + # GH6505 + other_names = set() + for var in names: + maybe_midx = self._indexes.get(var, None) + if isinstance(maybe_midx, PandasMultiIndex): + idx_coord_names = set(maybe_midx.index.names + [maybe_midx.dim]) + idx_other_names = idx_coord_names - set(names) + other_names.update(idx_other_names) + if other_names: + names |= set(other_names) + warnings.warn( + f"Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. " + f"Please also drop the following variables: {other_names!r} to avoid an error in the future.", + DeprecationWarning, + stacklevel=2, + ) + assert_no_index_corrupted(self.xindexes, names) variables = {k: v for k, v in self._variables.items() if k not in names} @@ -4545,7 +4742,9 @@ def drop_vars( variables, coord_names=coord_names, indexes=indexes ) - def drop(self, labels=None, dim=None, *, errors="raise", **labels_kwargs): + def drop( + self, labels=None, dim=None, *, errors: ErrorOptions = "raise", **labels_kwargs + ): """Backward compatible method based on `drop_vars` and `drop_sel` Using either `drop_vars` or `drop_sel` is encouraged @@ -4594,15 +4793,15 @@ def drop(self, labels=None, dim=None, *, errors="raise", **labels_kwargs): ) return self.drop_sel(labels, errors=errors) - def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs): + def drop_sel(self, labels=None, *, errors: ErrorOptions = "raise", **labels_kwargs): """Drop index labels from this dataset. Parameters ---------- labels : mapping of hashable to Any Index labels to drop - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the index labels passed are not in the dataset. If 'ignore', any given labels that are in the dataset are dropped and no error is raised. @@ -4729,7 +4928,10 @@ def drop_isel(self, indexers=None, **indexers_kwargs): return ds def drop_dims( - self, drop_dims: Hashable | Iterable[Hashable], *, errors: str = "raise" + self, + drop_dims: Hashable | Iterable[Hashable], + *, + errors: ErrorOptions = "raise", ) -> Dataset: """Drop dimensions and associated variables from this dataset. @@ -4769,7 +4971,7 @@ def drop_dims( def transpose( self, *dims: Hashable, - missing_dims: str = "raise", + missing_dims: ErrorOptionsWithWarn = "raise", ) -> Dataset: """Return a new Dataset object with all array dimensions transposed. @@ -4984,6 +5186,7 @@ def interpolate_na( provided. - 'barycentric', 'krog', 'pchip', 'spline', 'akima': use their respective :py:class:`scipy.interpolate` classes. + use_coordinate : bool, str, default: True Specifies which index to use as the x values in the interpolation formulated as `y = f(x)`. If False, values are treated as if @@ -5268,7 +5471,7 @@ def map( args: Iterable[Any] = (), **kwargs: Any, ) -> Dataset: - """Apply a function to each variable in this dataset + """Apply a function to each data variable in this dataset Parameters ---------- @@ -5470,7 +5673,7 @@ def to_array(self, dim="variable", name=None): return DataArray._construct_direct(variable, coords, name, indexes) def _normalize_dim_order( - self, dim_order: list[Hashable] = None + self, dim_order: Sequence[Hashable] | None = None ) -> dict[Hashable, int]: """ Check the validity of the provided dimensions if any and return the mapping @@ -5478,7 +5681,7 @@ def _normalize_dim_order( Parameters ---------- - dim_order + dim_order: Sequence of Hashable or None, optional Dimension order to validate (default to the alphabetical order if None). Returns @@ -5551,7 +5754,7 @@ def to_dataframe(self, dim_order: list[Hashable] = None) -> pd.DataFrame: Returns ------- - result + result : DataFrame Dataset as a pandas DataFrame. """ @@ -5777,7 +5980,7 @@ def to_dask_dataframe(self, dim_order=None, set_index=False): return df - def to_dict(self, data=True): + def to_dict(self, data: bool = True, encoding: bool = False) -> dict: """ Convert this dataset to a dictionary following xarray naming conventions. @@ -5791,21 +5994,34 @@ def to_dict(self, data=True): data : bool, optional Whether to include the actual data in the dictionary. When set to False, returns just the schema. + encoding : bool, optional + Whether to include the Dataset's encoding in the dictionary. + + Returns + ------- + d : dict See Also -------- Dataset.from_dict + DataArray.to_dict """ - d = { + d: dict = { "coords": {}, "attrs": decode_numpy_dict_values(self.attrs), "dims": dict(self.dims), "data_vars": {}, } for k in self.coords: - d["coords"].update({k: self[k].variable.to_dict(data=data)}) + d["coords"].update( + {k: self[k].variable.to_dict(data=data, encoding=encoding)} + ) for k in self.data_vars: - d["data_vars"].update({k: self[k].variable.to_dict(data=data)}) + d["data_vars"].update( + {k: self[k].variable.to_dict(data=data, encoding=encoding)} + ) + if encoding: + d["encoding"] = dict(self.encoding) return d @classmethod @@ -5894,6 +6110,7 @@ def from_dict(cls, d): obj = obj.set_coords(coords) obj.attrs.update(d.get("attrs", {})) + obj.encoding.update(d.get("encoding", {})) return obj @@ -6542,7 +6759,7 @@ def rank(self, dim, pct=False, keep_attrs=None): attrs = self.attrs if keep_attrs else None return self._replace(variables, coord_names, attrs=attrs) - def differentiate(self, coord, edge_order=1, datetime_unit=None): + def differentiate(self, coord, edge_order: Literal[1, 2] = 1, datetime_unit=None): """ Differentiate with the second order accurate central differences. @@ -7025,11 +7242,11 @@ def polyfit( self, dim: Hashable, deg: int, - skipna: bool = None, - rcond: float = None, + skipna: bool | None = None, + rcond: float | None = None, w: Hashable | Any = None, full: bool = False, - cov: bool | str = False, + cov: bool | Literal["unscaled"] = False, ): """ Least squares polynomial fit. @@ -7043,19 +7260,19 @@ def polyfit( Coordinate along which to fit the polynomials. deg : int Degree of the fitting polynomial. - skipna : bool, optional + skipna : bool or None, optional If True, removes all invalid values before fitting each 1D slices of the array. Default is True if data is stored in a dask.array or if there is any invalid values, False otherwise. - rcond : float, optional + rcond : float or None, optional Relative condition number to the fit. w : hashable or Any, optional Weights to apply to the y-coordinate of the sample points. Can be an array-like object or the name of a coordinate in the dataset. - full : bool, optional + full : bool, default: False Whether to return the residuals, matrix rank and singular values in addition to the coefficients. - cov : bool or str, optional + cov : bool or "unscaled", default: False Whether to return to the covariance matrix in addition to the coefficients. The matrix is not scaled if `cov='unscaled'`. @@ -7089,6 +7306,8 @@ def polyfit( numpy.polyval xarray.polyval """ + from .dataarray import DataArray + variables = {} skipna_da = skipna @@ -7122,10 +7341,10 @@ def polyfit( rank = np.linalg.matrix_rank(lhs) if full: - rank = xr.DataArray(rank, name=xname + "matrix_rank") + rank = DataArray(rank, name=xname + "matrix_rank") variables[rank.name] = rank _sing = np.linalg.svd(lhs, compute_uv=False) - sing = xr.DataArray( + sing = DataArray( _sing, dims=(degree_dim,), coords={degree_dim: np.arange(rank - 1, -1, -1)}, @@ -7178,7 +7397,7 @@ def polyfit( # Thus a ReprObject => polyfit was called on a DataArray name = "" - coeffs = xr.DataArray( + coeffs = DataArray( coeffs / scale_da, dims=[degree_dim] + list(stacked_coords.keys()), coords={degree_dim: np.arange(order)[::-1], **stacked_coords}, @@ -7189,7 +7408,7 @@ def polyfit( variables[coeffs.name] = coeffs if full or (cov is True): - residuals = xr.DataArray( + residuals = DataArray( residuals if dims_to_stack else residuals.squeeze(), dims=list(stacked_coords.keys()), coords=stacked_coords, @@ -7210,7 +7429,7 @@ def polyfit( "The number of data points must exceed order to scale the covariance matrix." ) fac = residuals / (x.shape[0] - order) - covariance = xr.DataArray(Vbase, dims=("cov_i", "cov_j")) * fac + covariance = DataArray(Vbase, dims=("cov_i", "cov_j")) * fac variables[name + "polyfit_covariance"] = covariance return Dataset(data_vars=variables, attrs=self.attrs.copy()) @@ -7218,16 +7437,16 @@ def polyfit( def pad( self, pad_width: Mapping[Any, int | tuple[int, int]] = None, - mode: str = "constant", + mode: PadModeOptions = "constant", stat_length: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, constant_values: ( - int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None + float | tuple[float, float] | Mapping[Any, tuple[float, float]] | None ) = None, end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, - reflect_type: str = None, + reflect_type: PadReflectOptions = None, **pad_width_kwargs: Any, ) -> Dataset: """Pad this dataset along one or more dimensions. @@ -7246,26 +7465,27 @@ def pad( Mapping with the form of {dim: (pad_before, pad_after)} describing the number of values padded along each dimension. {dim: pad} is a shortcut for pad_before = pad_after = pad - mode : str, default: "constant" - One of the following string values (taken from numpy docs). + mode : {"constant", "edge", "linear_ramp", "maximum", "mean", "median", \ + "minimum", "reflect", "symmetric", "wrap"}, default: "constant" + How to pad the DataArray (taken from numpy docs): - - constant: Pads with a constant value. - - edge: Pads with the edge values of array. - - linear_ramp: Pads with the linear ramp between end_value and the + - "constant": Pads with a constant value. + - "edge": Pads with the edge values of array. + - "linear_ramp": Pads with the linear ramp between end_value and the array edge value. - - maximum: Pads with the maximum value of all or part of the + - "maximum": Pads with the maximum value of all or part of the vector along each axis. - - mean: Pads with the mean value of all or part of the + - "mean": Pads with the mean value of all or part of the vector along each axis. - - median: Pads with the median value of all or part of the + - "median": Pads with the median value of all or part of the vector along each axis. - - minimum: Pads with the minimum value of all or part of the + - "minimum": Pads with the minimum value of all or part of the vector along each axis. - - reflect: Pads with the reflection of the vector mirrored on + - "reflect": Pads with the reflection of the vector mirrored on the first and last values of the vector along each axis. - - symmetric: Pads with the reflection of the vector mirrored + - "symmetric": Pads with the reflection of the vector mirrored along the edge of the array. - - wrap: Pads with the wrap of the vector along the axis. + - "wrap": Pads with the wrap of the vector along the axis. The first values are used to pad the end and the end values are used to pad the beginning. @@ -7299,7 +7519,7 @@ def pad( ``(constant,)`` or ``constant`` is a shortcut for ``before = after = constant`` for all axes. Default is 0. - reflect_type : {"even", "odd"}, optional + reflect_type : {"even", "odd", None}, optional Used in "reflect", and "symmetric". The "even" style is the default with an unaltered reflection around the edge value. For the "odd" style, the extended part of the array is created by @@ -7700,10 +7920,10 @@ def argmax(self, dim=None, **kwargs): def query( self, - queries: Mapping[Any, Any] = None, - parser: str = "pandas", - engine: str = None, - missing_dims: str = "raise", + queries: Mapping[Any, Any] | None = None, + parser: QueryParserOptions = "pandas", + engine: QueryEngineOptions = None, + missing_dims: ErrorOptionsWithWarn = "raise", **queries_kwargs: Any, ) -> Dataset: """Return a new dataset with each array indexed along the specified @@ -7713,8 +7933,8 @@ def query( Parameters ---------- - queries : dict, optional - A dict with keys matching dimensions and values given by strings + queries : dict-like, optional + A dict-like with keys matching dimensions and values given by strings containing Python expressions to be evaluated against the data variables in the dataset. The expressions will be evaluated using the pandas eval() function, and can contain any valid Python expressions but cannot @@ -7736,7 +7956,7 @@ def query( Dataset: - "raise": raise an exception - - "warning": raise a warning, and ignore the missing dimensions + - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions **queries_kwargs : {dim: query, ...}, optional @@ -7863,6 +8083,10 @@ def curvefit( """ from scipy.optimize import curve_fit + from .alignment import broadcast + from .computation import apply_ufunc + from .dataarray import _THIS_ARRAY, DataArray + if p0 is None: p0 = {} if bounds is None: @@ -7879,7 +8103,7 @@ def curvefit( if ( isinstance(coords, str) - or isinstance(coords, xr.DataArray) + or isinstance(coords, DataArray) or not isinstance(coords, Iterable) ): coords = [coords] @@ -7898,7 +8122,7 @@ def curvefit( ) # Broadcast all coords with each other - coords_ = xr.broadcast(*coords_) + coords_ = broadcast(*coords_) coords_ = [ coord.broadcast_like(self, exclude=preserved_dims) for coord in coords_ ] @@ -7933,14 +8157,14 @@ def _wrapper(Y, *coords_, **kwargs): popt, pcov = curve_fit(func, x, y, **kwargs) return popt, pcov - result = xr.Dataset() + result = Dataset() for name, da in self.data_vars.items(): - if name is xr.core.dataarray._THIS_ARRAY: + if name is _THIS_ARRAY: name = "" else: name = f"{str(name)}_" - popt, pcov = xr.apply_ufunc( + popt, pcov = apply_ufunc( _wrapper, da, *coords_, @@ -7971,7 +8195,7 @@ def _wrapper(Y, *coords_, **kwargs): def drop_duplicates( self, - dim: Hashable | Iterable[Hashable] | ..., + dim: Hashable | Iterable[Hashable], keep: Literal["first", "last"] | Literal[False] = "first", ): """Returns a new Dataset with duplicate dimension values removed. @@ -7995,9 +8219,11 @@ def drop_duplicates( DataArray.drop_duplicates """ if isinstance(dim, str): - dims = (dim,) + dims: Iterable = (dim,) elif dim is ...: dims = self.dims + elif not isinstance(dim, Iterable): + dims = [dim] else: dims = dim diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index b85d0e1645e..5f09abdbcfd 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -30,7 +30,7 @@ import dask.array as dask_array from dask.base import tokenize except ImportError: - dask_array = None + dask_array = None # type: ignore def _dask_or_eager_func( @@ -431,7 +431,14 @@ def datetime_to_numeric(array, offset=None, datetime_unit=None, dtype=float): # Compute timedelta object. # For np.datetime64, this can silently yield garbage due to overflow. # One option is to enforce 1970-01-01 as the universal offset. - array = array - offset + + # This map_blocks call is for backwards compatibility. + # dask == 2021.04.1 does not support subtracting object arrays + # which is required for cftime + if is_duck_dask_array(array) and np.issubdtype(array.dtype, object): + array = array.map_blocks(lambda a, b: a - b, offset, meta=array._meta) + else: + array = array - offset # Scalar is converted to 0d-array if not hasattr(array, "dtype"): @@ -517,10 +524,19 @@ def pd_timedelta_to_float(value, datetime_unit): return np_timedelta64_to_float(value, datetime_unit) +def _timedelta_to_seconds(array): + return np.reshape([a.total_seconds() for a in array.ravel()], array.shape) * 1e6 + + def py_timedelta_to_float(array, datetime_unit): """Convert a timedelta object to a float, possibly at a loss of resolution.""" - array = np.asarray(array) - array = np.reshape([a.total_seconds() for a in array.ravel()], array.shape) * 1e6 + array = asarray(array) + if is_duck_dask_array(array): + array = array.map_blocks( + _timedelta_to_seconds, meta=np.array([], dtype=np.float64) + ) + else: + array = _timedelta_to_seconds(array) conversion_factor = np.timedelta64(1, "us") / np.timedelta64(1, datetime_unit) return conversion_factor * array diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 26a2bbd4170..e8e24c5b066 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -326,7 +326,7 @@ def __init__( if not hashable(group): raise TypeError( "`group` must be an xarray.DataArray or the " - "name of an xarray variable or dimension." + "name of an xarray variable or dimension. " f"Received {group!r} instead." ) group = obj[group] @@ -355,6 +355,7 @@ def __init__( if bins is not None: if duck_array_ops.isnull(bins).all(): raise ValueError("All bin edges are NaN.") + if isinstance(bins, int): _, bins = pd.cut(group.values, bins, **cut_kwargs, retbins=True) @@ -457,6 +458,10 @@ def _initialize_old(self): self._unique_coord = unique_coord self._stacked_dim = stacked_dim self._inserted_dims = inserted_dims + self._full_index = full_index + self._restore_coord_dims = restore_coord_dims + self._bins = bins + self._squeeze = squeeze @property def _group_indices(self): @@ -636,60 +641,67 @@ def _maybe_unstack(self, obj): return obj def _flox_reduce(self, dim, **kwargs): + """Adaptor function that translates our groupby API to that of flox.""" from flox.xarray import xarray_reduce - from .dataarray import DataArray from .dataset import Dataset - # TODO: fix this - kwargs.pop("numeric_only", None) + obj = self._original_obj + + # preserve current strategy (approximately) for dask groupby. + # We want to control the default anyway to prevent surprises + # if flox decides to change its default + kwargs.setdefault("method", "split-reduce") + + numeric_only = kwargs.pop("numeric_only", None) + if numeric_only: + non_numeric = { + name: var + for name, var in obj.data_vars.items() + if not (np.issubdtype(var.dtype, np.number) or (var.dtype == np.bool_)) + } + else: + non_numeric = {} # weird backcompat # reducing along a unique indexed dimension with squeeze=True # should raise an error - if (dim is None or dim == self._group.name) and ( - self._bins is None and self._group.name in self._obj.xindexes - ): - # TODO: switch to xindexes after we can use is_unique - index = self._obj.indexes[self._group.name] + if ( + dim is None or dim == self._group.name + ) and self._group.name in obj.xindexes: + index = obj.indexes[self._group.name] if index.is_unique and self._squeeze: raise ValueError(f"cannot reduce over dimensions {self._group.name!r}") - # this creates a label DataArray since resample doesn't do that somehow - if self._grouper is not None: - repeats = [] - for slicer in self._group_indices: - stop = ( - slicer.stop - if slicer.stop is not None - else self._obj.sizes[self._group_dim] - ) - repeats.append(stop - slicer.start) - labels = np.repeat(self._unique_coord.data, repeats) - group = DataArray( - labels, dims=(self._group_dim,), name=self._unique_coord.name - ) - else: - if isinstance(self._original_group, _DummyGroup): - group = self._original_group.name + # group is only passed by resample + group = kwargs.pop("group", None) + if group is None: + if isinstance(self._unstacked_group, _DummyGroup): + group = self._unstacked_group.name else: - group = self._original_group + group = self._unstacked_group - # Do this so we raise the same error message whether flox is present or not. - # Better to control it here than in flox. + unindexed_dims = tuple() if isinstance(group, str): + if group in obj.dims and group not in obj._indexes and self._bins is None: + unindexed_dims = (group,) group = self._original_obj[group] - if dim not in (None, Ellipsis): - if isinstance(dim, str): - dim = (dim,) - if any( - d not in group.dims and d not in self._original_obj.dims for d in dim - ): - raise ValueError(f"cannot reduce over dimensions {dim}.") - - # TODO: handle bins=N in flox + + if isinstance(dim, str): + dim = (dim,) + elif dim is None: + dim = group.dims + elif dim is Ellipsis: + dim = tuple(self._original_obj.dims) + + # Do this so we raise the same error message whether flox is present or not. + # Better to control it here than in flox. + if any(d not in group.dims and d not in self._original_obj.dims for d in dim): + raise ValueError(f"cannot reduce over dimensions {dim}.") + if self._bins is not None: - expected_groups = (self._bins,) + # TODO: fix this; When binning by time, self._bins is a DatetimeIndex + expected_groups = (np.array(self._bins),) isbin = (True,) # This is an annoying hack. Xarray returns np.nan # when there are no observations in a bin, instead of 0. @@ -704,11 +716,11 @@ def _flox_reduce(self, dim, **kwargs): # flox's default would not set np.nan for integer dtypes kwargs.setdefault("fill_value", np.nan) else: - expected_groups = None # (self._unique_coord.values,) + expected_groups = None isbin = False result = xarray_reduce( - self._original_obj, + self._original_obj.drop_vars(non_numeric), group, dim=dim, expected_groups=expected_groups, @@ -716,24 +728,31 @@ def _flox_reduce(self, dim, **kwargs): **kwargs, ) + # Ignore error when the groupby reduction is effectively + # a reduction of the underlying dataset + result = result.drop_vars(unindexed_dims, errors="ignore") + + # broadcast and restore non-numeric data variables (backcompat) + for name, var in non_numeric.items(): + if all(d not in var.dims for d in dim): + result[name] = var.variable.set_dims( + (group.name,) + var.dims, (result.sizes[group.name],) + var.shape + ) + if self._bins is not None: - # bins provided to dask_groupby are at full precision + # bins provided to flox are at full precision # the bin edge labels have a default precision of 3 # reassign to fix that. - # new_coord = [ - # pd.Interval(inter.left, inter.right) for inter in self._full_index - # ] - # result[self._group.name] = new_coord + new_coord = [ + pd.Interval(inter.left, inter.right) for inter in self._full_index + ] + result[self._group.name] = new_coord # Fix dimension order when binning a dimension coordinate # Needed as long as we do a separate code path for pint; # For some reason Datasets and DataArrays behave differently! if isinstance(self._obj, Dataset) and self._group_dim in self._obj.dims: result = result.transpose(self._group.name, ...) - if self._grouper is not None: # self._unique_coord.name == "__resample_dim__": - result = self._maybe_restore_empty_groups(result) - # TODO: make this cleaner; the renaming happens in DatasetResample.map - result = result.rename(dict(__resample_dim__=self._group_dim)) return result def fillna(self, value): diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py index e02e1f569b2..ee3ef17ed65 100644 --- a/xarray/core/indexes.py +++ b/xarray/core/indexes.py @@ -22,10 +22,10 @@ from . import formatting, nputils, utils from .indexing import IndexSelResult, PandasIndexingAdapter, PandasMultiIndexingAdapter -from .types import T_Index from .utils import Frozen, get_valid_numpy_dtype, is_dict_like, is_scalar if TYPE_CHECKING: + from .types import ErrorOptions, T_Index from .variable import Variable IndexVars = Dict[Any, "Variable"] @@ -1098,7 +1098,7 @@ def is_multi(self, key: Hashable) -> bool: return len(self._id_coord_names[self._coord_name_id[key]]) > 1 def get_all_coords( - self, key: Hashable, errors: str = "raise" + self, key: Hashable, errors: ErrorOptions = "raise" ) -> dict[Hashable, Variable]: """Return all coordinates having the same index. @@ -1106,7 +1106,7 @@ def get_all_coords( ---------- key : hashable Index key. - errors : {"raise", "ignore"}, optional + errors : {"raise", "ignore"}, default: "raise" If "raise", raises a ValueError if `key` is not in indexes. If "ignore", an empty tuple is returned instead. @@ -1129,7 +1129,7 @@ def get_all_coords( return {k: self._variables[k] for k in all_coord_names} def get_all_dims( - self, key: Hashable, errors: str = "raise" + self, key: Hashable, errors: ErrorOptions = "raise" ) -> Mapping[Hashable, int]: """Return all dimensions shared by an index. @@ -1137,7 +1137,7 @@ def get_all_dims( ---------- key : hashable Index key. - errors : {"raise", "ignore"}, optional + errors : {"raise", "ignore"}, default: "raise" If "raise", raises a ValueError if `key` is not in indexes. If "ignore", an empty tuple is returned instead. diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 27bd4954bc4..cbbd507eeff 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import enum import functools import operator @@ -6,19 +8,7 @@ from dataclasses import dataclass, field from datetime import timedelta from html import escape -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - Hashable, - Iterable, - List, - Mapping, - Optional, - Tuple, - Union, -) +from typing import TYPE_CHECKING, Any, Callable, Hashable, Iterable, Mapping import numpy as np import pandas as pd @@ -59,12 +49,12 @@ class IndexSelResult: """ - dim_indexers: Dict[Any, Any] - indexes: Dict[Any, "Index"] = field(default_factory=dict) - variables: Dict[Any, "Variable"] = field(default_factory=dict) - drop_coords: List[Hashable] = field(default_factory=list) - drop_indexes: List[Hashable] = field(default_factory=list) - rename_dims: Dict[Any, Hashable] = field(default_factory=dict) + dim_indexers: dict[Any, Any] + indexes: dict[Any, Index] = field(default_factory=dict) + variables: dict[Any, Variable] = field(default_factory=dict) + drop_coords: list[Hashable] = field(default_factory=list) + drop_indexes: list[Hashable] = field(default_factory=list) + rename_dims: dict[Any, Hashable] = field(default_factory=dict) def as_tuple(self): """Unlike ``dataclasses.astuple``, return a shallow copy. @@ -82,7 +72,7 @@ def as_tuple(self): ) -def merge_sel_results(results: List[IndexSelResult]) -> IndexSelResult: +def merge_sel_results(results: list[IndexSelResult]) -> IndexSelResult: all_dims_count = Counter([dim for res in results for dim in res.dim_indexers]) duplicate_dims = {k: v for k, v in all_dims_count.items() if v > 1} @@ -124,13 +114,13 @@ def group_indexers_by_index( obj: T_Xarray, indexers: Mapping[Any, Any], options: Mapping[str, Any], -) -> List[Tuple["Index", Dict[Any, Any]]]: +) -> list[tuple[Index, dict[Any, Any]]]: """Returns a list of unique indexes and their corresponding indexers.""" unique_indexes = {} - grouped_indexers: Mapping[Union[int, None], Dict] = defaultdict(dict) + grouped_indexers: Mapping[int | None, dict] = defaultdict(dict) for key, label in indexers.items(): - index: "Index" = obj.xindexes.get(key, None) + index: Index = obj.xindexes.get(key, None) if index is not None: index_id = id(index) @@ -787,7 +777,7 @@ class IndexingSupport(enum.Enum): def explicit_indexing_adapter( key: ExplicitIndexer, - shape: Tuple[int, ...], + shape: tuple[int, ...], indexing_support: IndexingSupport, raw_indexing_method: Callable, ) -> Any: @@ -821,8 +811,8 @@ def explicit_indexing_adapter( def decompose_indexer( - indexer: ExplicitIndexer, shape: Tuple[int, ...], indexing_support: IndexingSupport -) -> Tuple[ExplicitIndexer, ExplicitIndexer]: + indexer: ExplicitIndexer, shape: tuple[int, ...], indexing_support: IndexingSupport +) -> tuple[ExplicitIndexer, ExplicitIndexer]: if isinstance(indexer, VectorizedIndexer): return _decompose_vectorized_indexer(indexer, shape, indexing_support) if isinstance(indexer, (BasicIndexer, OuterIndexer)): @@ -848,9 +838,9 @@ def _decompose_slice(key, size): def _decompose_vectorized_indexer( indexer: VectorizedIndexer, - shape: Tuple[int, ...], + shape: tuple[int, ...], indexing_support: IndexingSupport, -) -> Tuple[ExplicitIndexer, ExplicitIndexer]: +) -> tuple[ExplicitIndexer, ExplicitIndexer]: """ Decompose vectorized indexer to the successive two indexers, where the first indexer will be used to index backend arrays, while the second one @@ -929,10 +919,10 @@ def _decompose_vectorized_indexer( def _decompose_outer_indexer( - indexer: Union[BasicIndexer, OuterIndexer], - shape: Tuple[int, ...], + indexer: BasicIndexer | OuterIndexer, + shape: tuple[int, ...], indexing_support: IndexingSupport, -) -> Tuple[ExplicitIndexer, ExplicitIndexer]: +) -> tuple[ExplicitIndexer, ExplicitIndexer]: """ Decompose outer indexer to the successive two indexers, where the first indexer will be used to index backend arrays, while the second one @@ -973,7 +963,7 @@ def _decompose_outer_indexer( return indexer, BasicIndexer(()) assert isinstance(indexer, (OuterIndexer, BasicIndexer)) - backend_indexer: List[Any] = [] + backend_indexer: list[Any] = [] np_indexer = [] # make indexer positive pos_indexer: list[np.ndarray | int | np.number] = [] @@ -1395,7 +1385,7 @@ def __array__(self, dtype: DTypeLike = None) -> np.ndarray: return np.asarray(array.values, dtype=dtype) @property - def shape(self) -> Tuple[int]: + def shape(self) -> tuple[int]: return (len(self.array),) def _convert_scalar(self, item): @@ -1420,13 +1410,13 @@ def _convert_scalar(self, item): def __getitem__( self, indexer - ) -> Union[ - "PandasIndexingAdapter", - NumpyIndexingAdapter, - np.ndarray, - np.datetime64, - np.timedelta64, - ]: + ) -> ( + PandasIndexingAdapter + | NumpyIndexingAdapter + | np.ndarray + | np.datetime64 + | np.timedelta64 + ): key = indexer.tuple if isinstance(key, tuple) and len(key) == 1: # unpack key so it can index a pandas.Index object (pandas.Index @@ -1449,7 +1439,7 @@ def transpose(self, order) -> pd.Index: def __repr__(self) -> str: return f"{type(self).__name__}(array={self.array!r}, dtype={self.dtype!r})" - def copy(self, deep: bool = True) -> "PandasIndexingAdapter": + def copy(self, deep: bool = True) -> PandasIndexingAdapter: # Not the same as just writing `self.array.copy(deep=deep)`, as # shallow copies of the underlying numpy.ndarrays become deep ones # upon pickling @@ -1476,7 +1466,7 @@ def __init__( self, array: pd.MultiIndex, dtype: DTypeLike = None, - level: Optional[str] = None, + level: str | None = None, ): super().__init__(array, dtype) self.level = level @@ -1535,7 +1525,7 @@ def _repr_html_(self) -> str: array_repr = short_numpy_repr(self._get_array_subset()) return f"
{escape(array_repr)}
" - def copy(self, deep: bool = True) -> "PandasMultiIndexingAdapter": + def copy(self, deep: bool = True) -> PandasMultiIndexingAdapter: # see PandasIndexingAdapter.copy array = self.array.copy(deep=True) if deep else self.array return type(self)(array, self._dtype, self.level) diff --git a/xarray/core/merge.py b/xarray/core/merge.py index b428d4ae958..6262e031a2c 100644 --- a/xarray/core/merge.py +++ b/xarray/core/merge.py @@ -34,6 +34,7 @@ from .coordinates import Coordinates from .dataarray import DataArray from .dataset import Dataset + from .types import CombineAttrsOptions, CompatOptions, JoinOptions DimsLike = Union[Hashable, Sequence[Hashable]] ArrayLike = Any @@ -94,8 +95,8 @@ class MergeError(ValueError): def unique_variable( name: Hashable, variables: list[Variable], - compat: str = "broadcast_equals", - equals: bool = None, + compat: CompatOptions = "broadcast_equals", + equals: bool | None = None, ) -> Variable: """Return the unique variable from a list of variables or raise MergeError. @@ -207,9 +208,9 @@ def _assert_prioritized_valid( def merge_collected( grouped: dict[Hashable, list[MergeElement]], prioritized: Mapping[Any, MergeElement] = None, - compat: str = "minimal", - combine_attrs: str | None = "override", - equals: dict[Hashable, bool] = None, + compat: CompatOptions = "minimal", + combine_attrs: CombineAttrsOptions = "override", + equals: dict[Hashable, bool] | None = None, ) -> tuple[dict[Hashable, Variable], dict[Hashable, Index]]: """Merge dicts of variables, while resolving conflicts appropriately. @@ -219,6 +220,22 @@ def merge_collected( prioritized : mapping compat : str Type of equality check to use when checking for conflicts. + combine_attrs : {"drop", "identical", "no_conflicts", "drop_conflicts", \ + "override"} or callable, default: "override" + A callable or a string indicating how to combine attrs of the objects being + merged: + + - "drop": empty attrs on returned Dataset. + - "identical": all attrs must be the same on every object. + - "no_conflicts": attrs from all objects are combined, any that have + the same name must also have the same value. + - "drop_conflicts": attrs from all objects are combined, any that have + the same name but different values are dropped. + - "override": skip comparing and copy attrs from the first dataset to + the result. + + If a callable, it must expect a sequence of ``attrs`` dicts and a context object + as its only parameters. equals : mapping, optional corresponding to result of compat test @@ -376,7 +393,7 @@ def merge_coordinates_without_align( objects: list[Coordinates], prioritized: Mapping[Any, MergeElement] = None, exclude_dims: AbstractSet = frozenset(), - combine_attrs: str = "override", + combine_attrs: CombineAttrsOptions = "override", ) -> tuple[dict[Hashable, Variable], dict[Hashable, Index]]: """Merge variables/indexes from coordinates without automatic alignments. @@ -480,7 +497,9 @@ def coerce_pandas_values(objects: Iterable[CoercibleMapping]) -> list[DatasetLik def _get_priority_vars_and_indexes( - objects: list[DatasetLike], priority_arg: int | None, compat: str = "equals" + objects: list[DatasetLike], + priority_arg: int | None, + compat: CompatOptions = "equals", ) -> dict[Hashable, MergeElement]: """Extract the priority variable from a list of mappings. @@ -494,8 +513,19 @@ def _get_priority_vars_and_indexes( Dictionaries in which to find the priority variables. priority_arg : int or None Integer object whose variable should take priority. - compat : {"identical", "equals", "broadcast_equals", "no_conflicts"}, optional - Compatibility checks to use when merging variables. + compat : {"identical", "equals", "broadcast_equals", "no_conflicts", "override"}, optional + String indicating how to compare non-concatenated variables of the same name for + potential conflicts. This is passed down to merge. + + - "broadcast_equals": all values must be equal when variables are + broadcast against each other to ensure common dimensions. + - "equals": all values and dimensions must be the same. + - "identical": all values, dimensions and attributes must be the + same. + - "no_conflicts": only values which are not null in both datasets + must be equal. The returned dataset then contains the combination + of all non-null values. + - "override": skip comparing and pick variable from first dataset Returns ------- @@ -514,8 +544,8 @@ def _get_priority_vars_and_indexes( def merge_coords( objects: Iterable[CoercibleMapping], - compat: str = "minimal", - join: str = "outer", + compat: CompatOptions = "minimal", + join: JoinOptions = "outer", priority_arg: int | None = None, indexes: Mapping[Any, Index] | None = None, fill_value: object = dtypes.NA, @@ -665,9 +695,9 @@ class _MergeResult(NamedTuple): def merge_core( objects: Iterable[CoercibleMapping], - compat: str = "broadcast_equals", - join: str = "outer", - combine_attrs: str | None = "override", + compat: CompatOptions = "broadcast_equals", + join: JoinOptions = "outer", + combine_attrs: CombineAttrsOptions = "override", priority_arg: int | None = None, explicit_coords: Sequence | None = None, indexes: Mapping[Any, Any] | None = None, @@ -754,10 +784,10 @@ def merge_core( def merge( objects: Iterable[DataArray | CoercibleMapping], - compat: str = "no_conflicts", - join: str = "outer", + compat: CompatOptions = "no_conflicts", + join: JoinOptions = "outer", fill_value: object = dtypes.NA, - combine_attrs: str = "override", + combine_attrs: CombineAttrsOptions = "override", ) -> Dataset: """Merge any number of xarray objects into a single Dataset as variables. @@ -779,6 +809,7 @@ def merge( must be equal. The returned dataset then contains the combination of all non-null values. - "override": skip comparing and pick variable from first dataset + join : {"outer", "inner", "left", "right", "exact"}, optional String indicating how to combine differing indexes in objects. @@ -791,6 +822,7 @@ def merge( - "override": if indexes are of same size, rewrite indexes to be those of the first object with that dimension. Indexes for the same dimension must have the same size in all objects. + fill_value : scalar or dict-like, optional Value to use for newly missing values. If a dict-like, maps variable names to fill values. Use a data array's name to @@ -1002,10 +1034,10 @@ def dataset_merge_method( dataset: Dataset, other: CoercibleMapping, overwrite_vars: Hashable | Iterable[Hashable], - compat: str, - join: str, + compat: CompatOptions, + join: JoinOptions, fill_value: Any, - combine_attrs: str, + combine_attrs: CombineAttrsOptions, ) -> _MergeResult: """Guts of the Dataset.merge method.""" # we are locked into supporting overwrite_vars for the Dataset.merge diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 3d33631bebd..5e954c8ce27 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import datetime as dt import warnings from functools import partial from numbers import Number -from typing import Any, Callable, Dict, Hashable, Sequence, Union +from typing import TYPE_CHECKING, Any, Callable, Hashable, Sequence, get_args import numpy as np import pandas as pd @@ -14,11 +16,18 @@ from .duck_array_ops import datetime_to_numeric, push, timedelta_to_numeric from .options import OPTIONS, _get_keep_attrs from .pycompat import dask_version, is_duck_dask_array +from .types import Interp1dOptions, InterpOptions from .utils import OrderedSet, is_scalar from .variable import Variable, broadcast_variables +if TYPE_CHECKING: + from .dataarray import DataArray + from .dataset import Dataset + -def _get_nan_block_lengths(obj, dim: Hashable, index: Variable): +def _get_nan_block_lengths( + obj: Dataset | DataArray | Variable, dim: Hashable, index: Variable +): """ Return an object where each NaN element in 'obj' is replaced by the length of the gap the element is in. @@ -48,8 +57,8 @@ def _get_nan_block_lengths(obj, dim: Hashable, index: Variable): class BaseInterpolator: """Generic interpolator class for normalizing interpolation methods""" - cons_kwargs: Dict[str, Any] - call_kwargs: Dict[str, Any] + cons_kwargs: dict[str, Any] + call_kwargs: dict[str, Any] f: Callable method: str @@ -213,7 +222,7 @@ def _apply_over_vars_with_dim(func, self, dim=None, **kwargs): def get_clean_interp_index( - arr, dim: Hashable, use_coordinate: Union[str, bool] = True, strict: bool = True + arr, dim: Hashable, use_coordinate: str | bool = True, strict: bool = True ): """Return index to use for x values in interpolation or curve fitting. @@ -300,10 +309,10 @@ def get_clean_interp_index( def interp_na( self, dim: Hashable = None, - use_coordinate: Union[bool, str] = True, - method: str = "linear", + use_coordinate: bool | str = True, + method: InterpOptions = "linear", limit: int = None, - max_gap: Union[int, float, str, pd.Timedelta, np.timedelta64, dt.timedelta] = None, + max_gap: int | float | str | pd.Timedelta | np.timedelta64 | dt.timedelta = None, keep_attrs: bool = None, **kwargs, ): @@ -461,28 +470,20 @@ def _import_interpolant(interpolant, method): raise ImportError(f"Interpolation with method {method} requires scipy.") from e -def _get_interpolator(method, vectorizeable_only=False, **kwargs): +def _get_interpolator( + method: InterpOptions, vectorizeable_only: bool = False, **kwargs +): """helper function to select the appropriate interpolator class returns interpolator class and keyword arguments for the class """ - interp1d_methods = [ - "linear", - "nearest", - "zero", - "slinear", - "quadratic", - "cubic", - "polynomial", - ] - valid_methods = interp1d_methods + [ - "barycentric", - "krog", - "pchip", - "spline", - "akima", + interp_class: type[NumpyInterpolator] | type[ScipyInterpolator] | type[ + SplineInterpolator ] + interp1d_methods = get_args(Interp1dOptions) + valid_methods = tuple(vv for v in get_args(InterpOptions) for vv in get_args(v)) + # prioritize scipy.interpolate if ( method == "linear" @@ -589,7 +590,7 @@ def _floatize_x(x, new_x): return x, new_x -def interp(var, indexes_coords, method, **kwargs): +def interp(var, indexes_coords, method: InterpOptions, **kwargs): """Make an interpolation of Variable Parameters @@ -642,7 +643,7 @@ def interp(var, indexes_coords, method, **kwargs): result = Variable(new_dims, interped, attrs=var.attrs) # dimension of the output array - out_dims = OrderedSet() + out_dims: OrderedSet = OrderedSet() for d in var.dims: if d in dims: out_dims.update(indexes_coords[d][1].dims) @@ -652,7 +653,7 @@ def interp(var, indexes_coords, method, **kwargs): return result -def interp_func(var, x, new_x, method, kwargs): +def interp_func(var, x, new_x, method: InterpOptions, kwargs): """ multi-dimensional interpolation for array-like. Interpolated axes should be located in the last position. diff --git a/xarray/core/nanops.py b/xarray/core/nanops.py index c1a4d629f97..fa96bd6e150 100644 --- a/xarray/core/nanops.py +++ b/xarray/core/nanops.py @@ -11,7 +11,7 @@ from . import dask_array_compat except ImportError: - dask_array = None + dask_array = None # type: ignore[assignment] dask_array_compat = None # type: ignore[assignment] diff --git a/xarray/core/npcompat.py b/xarray/core/npcompat.py index b5b98052fe9..85a8f88aba6 100644 --- a/xarray/core/npcompat.py +++ b/xarray/core/npcompat.py @@ -28,7 +28,17 @@ # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from typing import TYPE_CHECKING, Any, Literal, Sequence, TypeVar, Union +from typing import ( + TYPE_CHECKING, + Any, + List, + Literal, + Sequence, + Tuple, + Type, + TypeVar, + Union, +) import numpy as np from packaging.version import Version @@ -36,6 +46,31 @@ # Type annotations stubs try: from numpy.typing import ArrayLike, DTypeLike + from numpy.typing._dtype_like import _DTypeLikeNested, _ShapeLike, _SupportsDType + + # Xarray requires a Mapping[Hashable, dtype] in many places which + # conflics with numpys own DTypeLike (with dtypes for fields). + # https://numpy.org/devdocs/reference/typing.html#numpy.typing.DTypeLike + # This is a copy of this DTypeLike that allows only non-Mapping dtypes. + DTypeLikeSave = Union[ + np.dtype, + # default data type (float64) + None, + # array-scalar types and generic types + Type[Any], + # character codes, type strings or comma-separated fields, e.g., 'float64' + str, + # (flexible_dtype, itemsize) + Tuple[_DTypeLikeNested, int], + # (fixed_dtype, shape) + Tuple[_DTypeLikeNested, _ShapeLike], + # (base_dtype, new_dtype) + Tuple[_DTypeLikeNested, _DTypeLikeNested], + # because numpy does the same? + List[Any], + # anything with a dtype attribute + _SupportsDType[np.dtype], + ] except ImportError: # fall back for numpy < 1.20, ArrayLike adapted from numpy.typing._array_like from typing import Protocol @@ -46,8 +81,14 @@ class _SupportsArray(Protocol): def __array__(self) -> np.ndarray: ... + class _SupportsDTypeFallback(Protocol): + @property + def dtype(self) -> np.dtype: + ... + else: _SupportsArray = Any + _SupportsDTypeFallback = Any _T = TypeVar("_T") _NestedSequence = Union[ @@ -72,7 +113,16 @@ def __array__(self) -> np.ndarray: # with the same name (ArrayLike and DTypeLike from the try block) ArrayLike = _ArrayLikeFallback # type: ignore # fall back for numpy < 1.20 - DTypeLike = Union[np.dtype, str] # type: ignore[misc] + DTypeLikeSave = Union[ # type: ignore[misc] + np.dtype, + str, + None, + Type[Any], + Tuple[Any, Any], + List[Any], + _SupportsDTypeFallback, + ] + DTypeLike = DTypeLikeSave # type: ignore[misc] if Version(np.__version__) >= Version("1.20.0"): diff --git a/xarray/core/resample.py b/xarray/core/resample.py index ed665ad4048..bcc4bfb90cd 100644 --- a/xarray/core/resample.py +++ b/xarray/core/resample.py @@ -1,13 +1,15 @@ import warnings from typing import Any, Callable, Hashable, Sequence, Union +import numpy as np + from ._reductions import DataArrayResampleReductions, DatasetResampleReductions -from .groupby import DataArrayGroupByBase, DatasetGroupByBase +from .groupby import DataArrayGroupByBase, DatasetGroupByBase, GroupBy RESAMPLE_DIM = "__resample_dim__" -class Resample: +class Resample(GroupBy): """An object that extends the `GroupBy` object with additional logic for handling specialized re-sampling operations. @@ -21,6 +23,29 @@ class Resample: """ + def _flox_reduce(self, dim, **kwargs): + + from .dataarray import DataArray + + kwargs.setdefault("method", "cohorts") + + # now create a label DataArray since resample doesn't do that somehow + repeats = [] + for slicer in self._group_indices: + stop = ( + slicer.stop + if slicer.stop is not None + else self._obj.sizes[self._group_dim] + ) + repeats.append(stop - slicer.start) + labels = np.repeat(self._unique_coord.data, repeats) + group = DataArray(labels, dims=(self._group_dim,), name=self._unique_coord.name) + + result = super()._flox_reduce(dim=dim, group=group, **kwargs) + result = self._maybe_restore_empty_groups(result) + result = result.rename({RESAMPLE_DIM: self._group_dim}) + return result + def _upsample(self, method, *args, **kwargs): """Dispatch function to call appropriate up-sampling methods on data. @@ -158,7 +183,7 @@ def _interpolate(self, kind="linear"): ) -class DataArrayResample(DataArrayGroupByBase, DataArrayResampleReductions, Resample): +class DataArrayResample(Resample, DataArrayGroupByBase, DataArrayResampleReductions): """DataArrayGroupBy object specialized to time resampling operations over a specified dimension """ @@ -249,7 +274,7 @@ def apply(self, func, args=(), shortcut=None, **kwargs): return self.map(func=func, shortcut=shortcut, args=args, **kwargs) -class DatasetResample(DatasetGroupByBase, DatasetResampleReductions, Resample): +class DatasetResample(Resample, DatasetGroupByBase, DatasetResampleReductions): """DatasetGroupBy object specialized to resampling a specified dimension""" def __init__(self, *args, dim=None, resample_dim=None, **kwargs): diff --git a/xarray/core/types.py b/xarray/core/types.py index 3f368501b25..f4f86bafc93 100644 --- a/xarray/core/types.py +++ b/xarray/core/types.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, TypeVar, Union +from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence, TypeVar, Union import numpy as np @@ -16,7 +16,7 @@ try: from dask.array import Array as DaskArray except ImportError: - DaskArray = np.ndarray + DaskArray = np.ndarray # type: ignore T_Dataset = TypeVar("T_Dataset", bound="Dataset") @@ -24,6 +24,8 @@ T_Variable = TypeVar("T_Variable", bound="Variable") T_Index = TypeVar("T_Index", bound="Index") +T_DataArrayOrSet = TypeVar("T_DataArrayOrSet", bound=Union["Dataset", "DataArray"]) + # Maybe we rename this to T_Data or something less Fortran-y? T_Xarray = TypeVar("T_Xarray", "DataArray", "Dataset") T_DataWithCoords = TypeVar("T_DataWithCoords", bound="DataWithCoords") @@ -33,3 +35,67 @@ DaCompatible = Union["DataArray", "Variable", "DataArrayGroupBy", "ScalarOrArray"] VarCompatible = Union["Variable", "ScalarOrArray"] GroupByIncompatible = Union["Variable", "GroupBy"] + +ErrorOptions = Literal["raise", "ignore"] +ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"] + +CompatOptions = Literal[ + "identical", "equals", "broadcast_equals", "no_conflicts", "override", "minimal" +] +ConcatOptions = Literal["all", "minimal", "different"] +CombineAttrsOptions = Union[ + Literal["drop", "identical", "no_conflicts", "drop_conflicts", "override"], + Callable[..., Any], +] +JoinOptions = Literal["outer", "inner", "left", "right", "exact", "override"] + +Interp1dOptions = Literal[ + "linear", "nearest", "zero", "slinear", "quadratic", "cubic", "polynomial" +] +InterpolantOptions = Literal["barycentric", "krog", "pchip", "spline", "akima"] +InterpOptions = Union[Interp1dOptions, InterpolantOptions] + +DatetimeUnitOptions = Literal[ + "Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns", "ps", "fs", "as" +] + +QueryEngineOptions = Literal["python", "numexpr", None] +QueryParserOptions = Literal["pandas", "python"] + +ReindexMethodOptions = Literal["nearest", "pad", "ffill", "backfill", "bfill", None] + +PadModeOptions = Literal[ + "constant", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", +] +PadReflectOptions = Literal["even", "odd", None] + +CFCalendar = Literal[ + "standard", + "gregorian", + "proleptic_gregorian", + "noleap", + "365_day", + "360_day", + "julian", + "all_leap", + "366_day", +] + +# TODO: Wait until mypy supports recursive objects in combination with typevars +_T = TypeVar("_T") +NestedSequence = Union[ + _T, + Sequence[_T], + Sequence[Sequence[_T]], + Sequence[Sequence[Sequence[_T]]], + Sequence[Sequence[Sequence[Sequence[_T]]]], +] diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 9f5f25c5895..121710d19aa 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -16,6 +16,7 @@ Callable, Collection, Container, + Generic, Hashable, Iterable, Iterator, @@ -24,11 +25,15 @@ MutableSet, TypeVar, cast, + overload, ) import numpy as np import pandas as pd +if TYPE_CHECKING: + from .types import ErrorOptionsWithWarn + K = TypeVar("K") V = TypeVar("V") T = TypeVar("T") @@ -237,7 +242,8 @@ def remove_incompatible_items( del first_dict[k] -def is_dict_like(value: Any) -> bool: +# It's probably OK to give this as a TypeGuard; though it's not perfectly robust. +def is_dict_like(value: Any) -> TypeGuard[dict]: return hasattr(value, "keys") and hasattr(value, "__getitem__") @@ -266,7 +272,7 @@ def either_dict_or_kwargs( kw_kwargs: Mapping[str, T], func_name: str, ) -> Mapping[Hashable, T]: - if pos_kwargs is None: + if pos_kwargs is None or pos_kwargs == {}: # Need an explicit cast to appease mypy due to invariance; see # https://github.com/python/mypy/issues/6228 return cast(Mapping[Hashable, T], kw_kwargs) @@ -755,7 +761,9 @@ def __len__(self) -> int: def infix_dims( - dims_supplied: Collection, dims_all: Collection, missing_dims: str = "raise" + dims_supplied: Collection, + dims_all: Collection, + missing_dims: ErrorOptionsWithWarn = "raise", ) -> Iterator: """ Resolves a supplied list containing an ellipsis representing other items, to @@ -803,7 +811,7 @@ def get_temp_dimname(dims: Container[Hashable], new_dim: Hashable) -> Hashable: def drop_dims_from_indexers( indexers: Mapping[Any, Any], dims: list | Mapping[Any, int], - missing_dims: str, + missing_dims: ErrorOptionsWithWarn, ) -> Mapping[Hashable, Any]: """Depending on the setting of missing_dims, drop any dimensions from indexers that are not present in dims. @@ -849,7 +857,7 @@ def drop_dims_from_indexers( def drop_missing_dims( - supplied_dims: Collection, dims: Collection, missing_dims: str + supplied_dims: Collection, dims: Collection, missing_dims: ErrorOptionsWithWarn ) -> Collection: """Depending on the setting of missing_dims, drop any dimensions from supplied_dims that are not present in dims. @@ -890,7 +898,10 @@ def drop_missing_dims( ) -class UncachedAccessor: +_Accessor = TypeVar("_Accessor") + + +class UncachedAccessor(Generic[_Accessor]): """Acts like a property, but on both classes and class instances This class is necessary because some tools (e.g. pydoc and sphinx) @@ -898,14 +909,22 @@ class UncachedAccessor: accessor. """ - def __init__(self, accessor): + def __init__(self, accessor: type[_Accessor]) -> None: self._accessor = accessor - def __get__(self, obj, cls): + @overload + def __get__(self, obj: None, cls) -> type[_Accessor]: + ... + + @overload + def __get__(self, obj: object, cls) -> _Accessor: + ... + + def __get__(self, obj: None | object, cls) -> type[_Accessor] | _Accessor: if obj is None: return self._accessor - return self._accessor(obj) + return self._accessor(obj) # type: ignore # assume it is a valid accessor! # Singleton type, as per https://github.com/python/typing/pull/240 diff --git a/xarray/core/variable.py b/xarray/core/variable.py index a21cf8c2d97..798e5a9f43e 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -5,7 +5,7 @@ import numbers import warnings from datetime import timedelta -from typing import TYPE_CHECKING, Any, Hashable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Hashable, Literal, Mapping, Sequence import numpy as np import pandas as pd @@ -59,7 +59,12 @@ BASIC_INDEXING_TYPES = integer_types + (slice,) if TYPE_CHECKING: - from .types import T_Variable + from .types import ( + ErrorOptionsWithWarn, + PadModeOptions, + PadReflectOptions, + T_Variable, + ) class MissingDimensionsError(ValueError): @@ -533,13 +538,17 @@ def to_index(self): """Convert this variable to a pandas.Index""" return self.to_index_variable().to_index() - def to_dict(self, data=True): + def to_dict(self, data: bool = True, encoding: bool = False) -> dict: """Dictionary representation of variable.""" item = {"dims": self.dims, "attrs": decode_numpy_dict_values(self.attrs)} if data: item["data"] = ensure_us_time_resolution(self.values).tolist() else: item.update({"dtype": str(self.dtype), "shape": self.shape}) + + if encoding: + item["encoding"] = dict(self.encoding) + return item @property @@ -1012,7 +1021,20 @@ def chunksizes(self) -> Mapping[Any, tuple[int, ...]]: _array_counter = itertools.count() - def chunk(self, chunks={}, name=None, lock=False): + def chunk( + self, + chunks: ( + int + | Literal["auto"] + | tuple[int, ...] + | tuple[tuple[int, ...], ...] + | Mapping[Any, None | int | tuple[int, ...]] + ) = {}, + name: str = None, + lock: bool = False, + inline_array: bool = False, + **chunks_kwargs: Any, + ) -> Variable: """Coerce this array's data into a dask array with the given chunks. If this variable is a non-dask array, it will be converted to dask @@ -1034,10 +1056,23 @@ def chunk(self, chunks={}, name=None, lock=False): lock : optional Passed on to :py:func:`dask.array.from_array`, if the array is not already as dask array. + inline_array: optional + Passed on to :py:func:`dask.array.from_array`, if the array is not + already as dask array. + **chunks_kwargs : {dim: chunks, ...}, optional + The keyword arguments form of ``chunks``. + One of chunks or chunks_kwargs must be provided. Returns ------- chunked : xarray.Variable + + See Also + -------- + Variable.chunks + Variable.chunksizes + xarray.unify_chunks + dask.array.from_array """ import dask.array as da @@ -1049,6 +1084,11 @@ def chunk(self, chunks={}, name=None, lock=False): ) chunks = {} + if isinstance(chunks, (float, str, int, tuple, list)): + pass # dask.array.from_array can handle these directly + else: + chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk") + if utils.is_dict_like(chunks): chunks = {self.get_axis_num(dim): chunk for dim, chunk in chunks.items()} @@ -1078,7 +1118,9 @@ def chunk(self, chunks={}, name=None, lock=False): if utils.is_dict_like(chunks): chunks = tuple(chunks.get(n, s) for n, s in enumerate(self.shape)) - data = da.from_array(data, chunks, name=name, lock=lock, **kwargs) + data = da.from_array( + data, chunks, name=name, lock=lock, inline_array=inline_array, **kwargs + ) return self._replace(data=data) @@ -1139,7 +1181,7 @@ def _to_dense(self): def isel( self: T_Variable, indexers: Mapping[Any, Any] = None, - missing_dims: str = "raise", + missing_dims: ErrorOptionsWithWarn = "raise", **indexers_kwargs: Any, ) -> T_Variable: """Return a new array indexed along the specified dimension(s). @@ -1153,7 +1195,7 @@ def isel( What to do if dimensions that should be selected from are not present in the DataArray: - "raise": raise an exception - - "warning": raise a warning, and ignore the missing dimensions + - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions Returns @@ -1272,15 +1314,17 @@ def _pad_options_dim_to_index( def pad( self, pad_width: Mapping[Any, int | tuple[int, int]] | None = None, - mode: str = "constant", + mode: PadModeOptions = "constant", stat_length: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, - constant_values: (int | tuple[int, int] | Mapping[Any, tuple[int, int]]) + constant_values: float + | tuple[float, float] + | Mapping[Any, tuple[float, float]] | None = None, end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, - reflect_type: str | None = None, + reflect_type: PadReflectOptions = None, **pad_width_kwargs: Any, ): """ @@ -1347,7 +1391,7 @@ def pad( pad_width_by_index = self._pad_options_dim_to_index(pad_width) # create pad_options_kwargs, numpy/dask requires only relevant kwargs to be nonempty - pad_option_kwargs = {} + pad_option_kwargs: dict[str, Any] = {} if stat_length is not None: pad_option_kwargs["stat_length"] = stat_length if constant_values is not None: @@ -1355,7 +1399,7 @@ def pad( if end_values is not None: pad_option_kwargs["end_values"] = end_values if reflect_type is not None: - pad_option_kwargs["reflect_type"] = reflect_type # type: ignore[assignment] + pad_option_kwargs["reflect_type"] = reflect_type array = np.pad( # type: ignore[call-overload] self.data.astype(dtype, copy=False), @@ -1415,14 +1459,14 @@ def roll(self, shifts=None, **shifts_kwargs): def transpose( self, - *dims, - missing_dims: str = "raise", + *dims: Hashable, + missing_dims: ErrorOptionsWithWarn = "raise", ) -> Variable: """Return a new Variable object with transposed dimensions. Parameters ---------- - *dims : str, optional + *dims : Hashable, optional By default, reverse the dimensions. Otherwise, reorder the dimensions to this order. missing_dims : {"raise", "warn", "ignore"}, default: "raise" @@ -2690,7 +2734,7 @@ def values(self, values): f"Please use DataArray.assign_coords, Dataset.assign_coords or Dataset.assign as appropriate." ) - def chunk(self, chunks={}, name=None, lock=False): + def chunk(self, chunks={}, name=None, lock=False, inline_array=False): # Dummy - do not chunk. This method is invoked e.g. by Dataset.chunk() return self.copy(deep=False) diff --git a/xarray/static/css/style.css b/xarray/static/css/style.css index b3b8a162e9a..9fa27c03359 100644 --- a/xarray/static/css/style.css +++ b/xarray/static/css/style.css @@ -14,6 +14,7 @@ } html[theme=dark], +body[data-theme=dark], body.vscode-dark { --xr-font-color0: rgba(255, 255, 255, 1); --xr-font-color2: rgba(255, 255, 255, 0.54); diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 65f0bc08261..ff477a40891 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -1,7 +1,10 @@ +from __future__ import annotations + import importlib import platform import warnings from contextlib import contextmanager, nullcontext +from typing import Any from unittest import mock # noqa: F401 import numpy as np @@ -40,7 +43,7 @@ ) -def _importorskip(modname, minversion=None): +def _importorskip(modname: str, minversion: str | None = None) -> tuple[bool, Any]: try: mod = importlib.import_module(modname) has = True diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 825c6f7130f..6f92b26b0c9 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -4,6 +4,7 @@ import math import os.path import pickle +import platform import re import shutil import sys @@ -81,7 +82,11 @@ _NON_STANDARD_CALENDARS, _STANDARD_CALENDARS, ) -from .test_dataset import create_append_test_data, create_test_data +from .test_dataset import ( + create_append_string_length_mismatch_test_data, + create_append_test_data, + create_test_data, +) try: import netCDF4 as nc4 @@ -1756,7 +1761,7 @@ def test_auto_chunk(self): assert v.chunks == original[k].chunks @requires_dask - @pytest.mark.filterwarnings("ignore:Specified Dask chunks") + @pytest.mark.filterwarnings("ignore:The specified Dask chunks separate") def test_manual_chunk(self): original = create_test_data().chunk({"dim1": 3, "dim2": 4, "dim3": 3}) @@ -2111,6 +2116,17 @@ def test_append_with_existing_encoding_raises(self): encoding={"da": {"compressor": None}}, ) + @pytest.mark.parametrize("dtype", ["U", "S"]) + def test_append_string_length_mismatch_raises(self, dtype): + ds, ds_to_append = create_append_string_length_mismatch_test_data(dtype) + with self.create_zarr_target() as store_target: + ds.to_zarr(store_target, mode="w") + with pytest.raises(ValueError, match="Mismatched dtypes for variable"): + ds_to_append.to_zarr( + store_target, + append_dim="time", + ) + def test_check_encoding_is_consistent_after_append(self): ds, ds_to_append, _ = create_append_test_data() @@ -2210,6 +2226,13 @@ def test_save_emptydim(self, chunk): with self.roundtrip(ds) as ds_reload: assert_identical(ds, ds_reload) + @requires_dask + def test_no_warning_from_open_emptydim_with_chunks(self): + ds = Dataset({"x": (("a", "b"), np.empty((5, 0)))}).chunk({"a": 1}) + with assert_no_warnings(): + with self.roundtrip(ds, open_kwargs=dict(chunks={"a": 1})) as ds_reload: + assert_identical(ds, ds_reload) + @pytest.mark.parametrize("consolidated", [False, True]) @pytest.mark.parametrize("compute", [False, True]) @pytest.mark.parametrize("use_dask", [False, True]) @@ -2420,6 +2443,22 @@ def test_write_read_select_write(self): with self.create_zarr_target() as final_store: ds_sel.to_zarr(final_store, mode="w") + @pytest.mark.parametrize("obj", [Dataset(), DataArray(name="foo")]) + def test_attributes(self, obj): + obj = obj.copy() + + obj.attrs["good"] = {"key": "value"} + ds = obj if isinstance(obj, Dataset) else obj.to_dataset() + with self.create_zarr_target() as store_target: + ds.to_zarr(store_target) + assert_identical(ds, xr.open_zarr(store_target)) + + obj.attrs["bad"] = DataArray() + ds = obj if isinstance(obj, Dataset) else obj.to_dataset() + with self.create_zarr_target() as store_target: + with pytest.raises(TypeError, match=r"Invalid attribute in Dataset.attrs."): + ds.to_zarr(store_target) + @requires_zarr class TestZarrDictStore(ZarrBase): @@ -3411,7 +3450,7 @@ def test_dataset_caching(self): actual.foo.values # no caching assert not actual.foo.variable._in_memory - def test_open_mfdataset(self): + def test_open_mfdataset(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: @@ -3434,14 +3473,14 @@ def test_open_mfdataset(self): open_mfdataset("http://some/remote/uri") @requires_fsspec - def test_open_mfdataset_no_files(self): + def test_open_mfdataset_no_files(self) -> None: pytest.importorskip("aiobotocore") # glob is attempted as of #4823, but finds no files with pytest.raises(OSError, match=r"no files"): open_mfdataset("http://some/remote/uri", engine="zarr") - def test_open_mfdataset_2d(self): + def test_open_mfdataset_2d(self) -> None: original = Dataset({"foo": (["x", "y"], np.random.randn(10, 8))}) with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: @@ -3470,7 +3509,7 @@ def test_open_mfdataset_2d(self): (2, 2, 2, 2), ) - def test_open_mfdataset_pathlib(self): + def test_open_mfdataset_pathlib(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: @@ -3483,7 +3522,7 @@ def test_open_mfdataset_pathlib(self): ) as actual: assert_identical(original, actual) - def test_open_mfdataset_2d_pathlib(self): + def test_open_mfdataset_2d_pathlib(self) -> None: original = Dataset({"foo": (["x", "y"], np.random.randn(10, 8))}) with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: @@ -3504,7 +3543,7 @@ def test_open_mfdataset_2d_pathlib(self): ) as actual: assert_identical(original, actual) - def test_open_mfdataset_2(self): + def test_open_mfdataset_2(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: @@ -3516,7 +3555,7 @@ def test_open_mfdataset_2(self): ) as actual: assert_identical(original, actual) - def test_attrs_mfdataset(self): + def test_attrs_mfdataset(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: @@ -3536,7 +3575,7 @@ def test_attrs_mfdataset(self): with pytest.raises(AttributeError, match=r"no attribute"): actual.test2 - def test_open_mfdataset_attrs_file(self): + def test_open_mfdataset_attrs_file(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_files(2) as (tmp1, tmp2): ds1 = original.isel(x=slice(5)) @@ -3553,7 +3592,7 @@ def test_open_mfdataset_attrs_file(self): # attributes from ds1 are not retained, e.g., assert "test1" not in actual.attrs - def test_open_mfdataset_attrs_file_path(self): + def test_open_mfdataset_attrs_file_path(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_files(2) as (tmp1, tmp2): tmp1 = Path(tmp1) @@ -3572,7 +3611,7 @@ def test_open_mfdataset_attrs_file_path(self): # attributes from ds1 are not retained, e.g., assert "test1" not in actual.attrs - def test_open_mfdataset_auto_combine(self): + def test_open_mfdataset_auto_combine(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10)), "x": np.arange(10)}) with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: @@ -3582,7 +3621,7 @@ def test_open_mfdataset_auto_combine(self): with open_mfdataset([tmp2, tmp1], combine="by_coords") as actual: assert_identical(original, actual) - def test_open_mfdataset_raise_on_bad_combine_args(self): + def test_open_mfdataset_raise_on_bad_combine_args(self) -> None: # Regression test for unhelpful error shown in #5230 original = Dataset({"foo": ("x", np.random.randn(10)), "x": np.arange(10)}) with create_tmp_file() as tmp1: @@ -3593,7 +3632,7 @@ def test_open_mfdataset_raise_on_bad_combine_args(self): open_mfdataset([tmp1, tmp2], concat_dim="x") @pytest.mark.xfail(reason="mfdataset loses encoding currently.") - def test_encoding_mfdataset(self): + def test_encoding_mfdataset(self) -> None: original = Dataset( { "foo": ("t", np.random.randn(10)), @@ -3615,7 +3654,7 @@ def test_encoding_mfdataset(self): assert actual.t.encoding["units"] == ds1.t.encoding["units"] assert actual.t.encoding["units"] != ds2.t.encoding["units"] - def test_preprocess_mfdataset(self): + def test_preprocess_mfdataset(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp: original.to_netcdf(tmp) @@ -3629,7 +3668,7 @@ def preprocess(ds): ) as actual: assert_identical(expected, actual) - def test_save_mfdataset_roundtrip(self): + def test_save_mfdataset_roundtrip(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) datasets = [original.isel(x=slice(5)), original.isel(x=slice(5, 10))] with create_tmp_file() as tmp1: @@ -3640,20 +3679,20 @@ def test_save_mfdataset_roundtrip(self): ) as actual: assert_identical(actual, original) - def test_save_mfdataset_invalid(self): + def test_save_mfdataset_invalid(self) -> None: ds = Dataset() with pytest.raises(ValueError, match=r"cannot use mode"): save_mfdataset([ds, ds], ["same", "same"]) with pytest.raises(ValueError, match=r"same length"): save_mfdataset([ds, ds], ["only one path"]) - def test_save_mfdataset_invalid_dataarray(self): + def test_save_mfdataset_invalid_dataarray(self) -> None: # regression test for GH1555 da = DataArray([1, 2]) with pytest.raises(TypeError, match=r"supports writing Dataset"): save_mfdataset([da], ["dataarray"]) - def test_save_mfdataset_pathlib_roundtrip(self): + def test_save_mfdataset_pathlib_roundtrip(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) datasets = [original.isel(x=slice(5)), original.isel(x=slice(5, 10))] with create_tmp_file() as tmp1: @@ -3666,7 +3705,7 @@ def test_save_mfdataset_pathlib_roundtrip(self): ) as actual: assert_identical(actual, original) - def test_open_and_do_math(self): + def test_open_and_do_math(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp: original.to_netcdf(tmp) @@ -3674,7 +3713,7 @@ def test_open_and_do_math(self): actual = 1.0 * ds assert_allclose(original, actual, decode_bytes=False) - def test_open_mfdataset_concat_dim_none(self): + def test_open_mfdataset_concat_dim_none(self) -> None: with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: data = Dataset({"x": 0}) @@ -3685,7 +3724,7 @@ def test_open_mfdataset_concat_dim_none(self): ) as actual: assert_identical(data, actual) - def test_open_mfdataset_concat_dim_default_none(self): + def test_open_mfdataset_concat_dim_default_none(self) -> None: with create_tmp_file() as tmp1: with create_tmp_file() as tmp2: data = Dataset({"x": 0}) @@ -3694,7 +3733,7 @@ def test_open_mfdataset_concat_dim_default_none(self): with open_mfdataset([tmp1, tmp2], combine="nested") as actual: assert_identical(data, actual) - def test_open_dataset(self): + def test_open_dataset(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp: original.to_netcdf(tmp) @@ -3708,7 +3747,7 @@ def test_open_dataset(self): assert isinstance(actual.foo.variable.data, np.ndarray) assert_identical(original, actual) - def test_open_single_dataset(self): + def test_open_single_dataset(self) -> None: # Test for issue GH #1988. This makes sure that the # concat_dim is utilized when specified in open_mfdataset(). rnddata = np.random.randn(10) @@ -3722,7 +3761,7 @@ def test_open_single_dataset(self): with open_mfdataset([tmp], concat_dim=dim, combine="nested") as actual: assert_identical(expected, actual) - def test_open_multi_dataset(self): + def test_open_multi_dataset(self) -> None: # Test for issue GH #1988 and #2647. This makes sure that the # concat_dim is utilized when specified in open_mfdataset(). # The additional wrinkle is to ensure that a length greater @@ -3747,7 +3786,7 @@ def test_open_multi_dataset(self): ) as actual: assert_identical(expected, actual) - def test_dask_roundtrip(self): + def test_dask_roundtrip(self) -> None: with create_tmp_file() as tmp: data = create_test_data() data.to_netcdf(tmp) @@ -3759,7 +3798,7 @@ def test_dask_roundtrip(self): with open_dataset(tmp2) as on_disk: assert_identical(data, on_disk) - def test_deterministic_names(self): + def test_deterministic_names(self) -> None: with create_tmp_file() as tmp: data = create_test_data() data.to_netcdf(tmp) @@ -3772,7 +3811,7 @@ def test_deterministic_names(self): assert dask_name[:13] == "open_dataset-" assert original_names == repeat_names - def test_dataarray_compute(self): + def test_dataarray_compute(self) -> None: # Test DataArray.compute() on dask backend. # The test for Dataset.compute() is already in DatasetIOBase; # however dask is the only tested backend which supports DataArrays @@ -3783,7 +3822,7 @@ def test_dataarray_compute(self): assert_allclose(actual, computed, decode_bytes=False) @pytest.mark.xfail - def test_save_mfdataset_compute_false_roundtrip(self): + def test_save_mfdataset_compute_false_roundtrip(self) -> None: from dask.delayed import Delayed original = Dataset({"foo": ("x", np.random.randn(10))}).chunk() @@ -3800,7 +3839,7 @@ def test_save_mfdataset_compute_false_roundtrip(self): ) as actual: assert_identical(actual, original) - def test_load_dataset(self): + def test_load_dataset(self) -> None: with create_tmp_file() as tmp: original = Dataset({"foo": ("x", np.random.randn(10))}) original.to_netcdf(tmp) @@ -3808,7 +3847,7 @@ def test_load_dataset(self): # this would fail if we used open_dataset instead of load_dataset ds.to_netcdf(tmp) - def test_load_dataarray(self): + def test_load_dataarray(self) -> None: with create_tmp_file() as tmp: original = Dataset({"foo": ("x", np.random.randn(10))}) original.to_netcdf(tmp) @@ -3817,6 +3856,27 @@ def test_load_dataarray(self): # load_dataarray ds.to_netcdf(tmp) + @pytest.mark.skipif( + ON_WINDOWS, + reason="counting number of tasks in graph fails on windows for some reason", + ) + def test_inline_array(self) -> None: + with create_tmp_file() as tmp: + original = Dataset({"foo": ("x", np.random.randn(10))}) + original.to_netcdf(tmp) + chunks = {"time": 10} + + def num_graph_nodes(obj): + return len(obj.__dask_graph__()) + + not_inlined_ds = open_dataset(tmp, inline_array=False, chunks=chunks) + inlined_ds = open_dataset(tmp, inline_array=True, chunks=chunks) + assert num_graph_nodes(inlined_ds) < num_graph_nodes(not_inlined_ds) + + not_inlined_da = open_dataarray(tmp, inline_array=False, chunks=chunks) + inlined_da = open_dataarray(tmp, inline_array=True, chunks=chunks) + assert num_graph_nodes(inlined_da) < num_graph_nodes(not_inlined_da) + @requires_scipy_or_netCDF4 @requires_pydap @@ -4296,7 +4356,7 @@ def create_tmp_geotiff( @requires_rasterio class TestRasterio: @requires_scipy_or_netCDF4 - def test_serialization(self): + def test_serialization(self) -> None: with create_tmp_geotiff(additional_attrs={}) as (tmp_file, expected): # Write it to a netcdf and read again (roundtrip) with pytest.warns(DeprecationWarning), xr.open_rasterio(tmp_file) as rioda: @@ -4902,7 +4962,7 @@ def new_dataset_and_coord_attrs(): @requires_scipy_or_netCDF4 class TestDataArrayToNetCDF: - def test_dataarray_to_netcdf_no_name(self): + def test_dataarray_to_netcdf_no_name(self) -> None: original_da = DataArray(np.arange(12).reshape((3, 4))) with create_tmp_file() as tmp: @@ -4911,7 +4971,7 @@ def test_dataarray_to_netcdf_no_name(self): with open_dataarray(tmp) as loaded_da: assert_identical(original_da, loaded_da) - def test_dataarray_to_netcdf_with_name(self): + def test_dataarray_to_netcdf_with_name(self) -> None: original_da = DataArray(np.arange(12).reshape((3, 4)), name="test") with create_tmp_file() as tmp: @@ -4920,7 +4980,7 @@ def test_dataarray_to_netcdf_with_name(self): with open_dataarray(tmp) as loaded_da: assert_identical(original_da, loaded_da) - def test_dataarray_to_netcdf_coord_name_clash(self): + def test_dataarray_to_netcdf_coord_name_clash(self) -> None: original_da = DataArray( np.arange(12).reshape((3, 4)), dims=["x", "y"], name="x" ) @@ -4931,7 +4991,7 @@ def test_dataarray_to_netcdf_coord_name_clash(self): with open_dataarray(tmp) as loaded_da: assert_identical(original_da, loaded_da) - def test_open_dataarray_options(self): + def test_open_dataarray_options(self) -> None: data = DataArray(np.arange(5), coords={"y": ("x", range(5))}, dims=["x"]) with create_tmp_file() as tmp: @@ -4942,13 +5002,13 @@ def test_open_dataarray_options(self): assert_identical(expected, loaded) @requires_scipy - def test_dataarray_to_netcdf_return_bytes(self): + def test_dataarray_to_netcdf_return_bytes(self) -> None: # regression test for GH1410 data = xr.DataArray([1, 2, 3]) output = data.to_netcdf() assert isinstance(output, bytes) - def test_dataarray_to_netcdf_no_name_pathlib(self): + def test_dataarray_to_netcdf_no_name_pathlib(self) -> None: original_da = DataArray(np.arange(12).reshape((3, 4))) with create_tmp_file() as tmp: @@ -4960,7 +5020,7 @@ def test_dataarray_to_netcdf_no_name_pathlib(self): @requires_scipy_or_netCDF4 -def test_no_warning_from_dask_effective_get(): +def test_no_warning_from_dask_effective_get() -> None: with create_tmp_file() as tmpfile: with assert_no_warnings(): ds = Dataset() @@ -4968,7 +5028,7 @@ def test_no_warning_from_dask_effective_get(): @requires_scipy_or_netCDF4 -def test_source_encoding_always_present(): +def test_source_encoding_always_present() -> None: # Test for GH issue #2550. rnddata = np.random.randn(10) original = Dataset({"foo": ("x", rnddata)}) @@ -4986,13 +5046,12 @@ def _assert_no_dates_out_of_range_warning(record): @requires_scipy_or_netCDF4 @pytest.mark.parametrize("calendar", _STANDARD_CALENDARS) -def test_use_cftime_standard_calendar_default_in_range(calendar): +def test_use_cftime_standard_calendar_default_in_range(calendar) -> None: x = [0, 1] time = [0, 720] units_date = "2000-01-01" units = "days since 2000-01-01" - original = DataArray(x, [("time", time)], name="x") - original = original.to_dataset() + original = DataArray(x, [("time", time)], name="x").to_dataset() for v in ["x", "time"]: original[v].attrs["units"] = units original[v].attrs["calendar"] = calendar @@ -5017,14 +5076,15 @@ def test_use_cftime_standard_calendar_default_in_range(calendar): @requires_scipy_or_netCDF4 @pytest.mark.parametrize("calendar", _STANDARD_CALENDARS) @pytest.mark.parametrize("units_year", [1500, 2500]) -def test_use_cftime_standard_calendar_default_out_of_range(calendar, units_year): +def test_use_cftime_standard_calendar_default_out_of_range( + calendar, units_year +) -> None: import cftime x = [0, 1] time = [0, 720] units = f"days since {units_year}-01-01" - original = DataArray(x, [("time", time)], name="x") - original = original.to_dataset() + original = DataArray(x, [("time", time)], name="x").to_dataset() for v in ["x", "time"]: original[v].attrs["units"] = units original[v].attrs["calendar"] = calendar @@ -5048,14 +5108,13 @@ def test_use_cftime_standard_calendar_default_out_of_range(calendar, units_year) @requires_scipy_or_netCDF4 @pytest.mark.parametrize("calendar", _ALL_CALENDARS) @pytest.mark.parametrize("units_year", [1500, 2000, 2500]) -def test_use_cftime_true(calendar, units_year): +def test_use_cftime_true(calendar, units_year) -> None: import cftime x = [0, 1] time = [0, 720] units = f"days since {units_year}-01-01" - original = DataArray(x, [("time", time)], name="x") - original = original.to_dataset() + original = DataArray(x, [("time", time)], name="x").to_dataset() for v in ["x", "time"]: original[v].attrs["units"] = units original[v].attrs["calendar"] = calendar @@ -5078,13 +5137,12 @@ def test_use_cftime_true(calendar, units_year): @requires_scipy_or_netCDF4 @pytest.mark.parametrize("calendar", _STANDARD_CALENDARS) -def test_use_cftime_false_standard_calendar_in_range(calendar): +def test_use_cftime_false_standard_calendar_in_range(calendar) -> None: x = [0, 1] time = [0, 720] units_date = "2000-01-01" units = "days since 2000-01-01" - original = DataArray(x, [("time", time)], name="x") - original = original.to_dataset() + original = DataArray(x, [("time", time)], name="x").to_dataset() for v in ["x", "time"]: original[v].attrs["units"] = units original[v].attrs["calendar"] = calendar @@ -5108,12 +5166,11 @@ def test_use_cftime_false_standard_calendar_in_range(calendar): @requires_scipy_or_netCDF4 @pytest.mark.parametrize("calendar", _STANDARD_CALENDARS) @pytest.mark.parametrize("units_year", [1500, 2500]) -def test_use_cftime_false_standard_calendar_out_of_range(calendar, units_year): +def test_use_cftime_false_standard_calendar_out_of_range(calendar, units_year) -> None: x = [0, 1] time = [0, 720] units = f"days since {units_year}-01-01" - original = DataArray(x, [("time", time)], name="x") - original = original.to_dataset() + original = DataArray(x, [("time", time)], name="x").to_dataset() for v in ["x", "time"]: original[v].attrs["units"] = units original[v].attrs["calendar"] = calendar @@ -5127,12 +5184,11 @@ def test_use_cftime_false_standard_calendar_out_of_range(calendar, units_year): @requires_scipy_or_netCDF4 @pytest.mark.parametrize("calendar", _NON_STANDARD_CALENDARS) @pytest.mark.parametrize("units_year", [1500, 2000, 2500]) -def test_use_cftime_false_nonstandard_calendar(calendar, units_year): +def test_use_cftime_false_nonstandard_calendar(calendar, units_year) -> None: x = [0, 1] time = [0, 720] units = f"days since {units_year}" - original = DataArray(x, [("time", time)], name="x") - original = original.to_dataset() + original = DataArray(x, [("time", time)], name="x").to_dataset() for v in ["x", "time"]: original[v].attrs["units"] = units original[v].attrs["calendar"] = calendar @@ -5200,7 +5256,7 @@ def test_extract_zarr_variable_encoding(): @requires_zarr @requires_fsspec @pytest.mark.filterwarnings("ignore:deallocating CachingFileManager") -def test_open_fsspec(): +def test_open_fsspec() -> None: import fsspec import zarr @@ -5242,7 +5298,7 @@ def test_open_fsspec(): @requires_h5netcdf @requires_netCDF4 -def test_load_single_value_h5netcdf(tmp_path): +def test_load_single_value_h5netcdf(tmp_path: Path) -> None: """Test that numeric single-element vector attributes are handled fine. At present (h5netcdf v0.8.1), the h5netcdf exposes single-valued numeric variable @@ -5267,7 +5323,7 @@ def test_load_single_value_h5netcdf(tmp_path): @pytest.mark.parametrize( "chunks", ["auto", -1, {}, {"x": "auto"}, {"x": -1}, {"x": "auto", "y": -1}] ) -def test_open_dataset_chunking_zarr(chunks, tmp_path): +def test_open_dataset_chunking_zarr(chunks, tmp_path: Path) -> None: encoded_chunks = 100 dask_arr = da.from_array( np.ones((500, 500), dtype="float64"), chunks=encoded_chunks @@ -5296,7 +5352,7 @@ def test_open_dataset_chunking_zarr(chunks, tmp_path): @pytest.mark.parametrize( "chunks", ["auto", -1, {}, {"x": "auto"}, {"x": -1}, {"x": "auto", "y": -1}] ) -@pytest.mark.filterwarnings("ignore:Specified Dask chunks") +@pytest.mark.filterwarnings("ignore:The specified Dask chunks separate") def test_chunking_consintency(chunks, tmp_path): encoded_chunks = {} dask_arr = da.from_array( @@ -5332,7 +5388,7 @@ def _check_guess_can_open_and_open(entrypoint, obj, engine, expected): @requires_netCDF4 -def test_netcdf4_entrypoint(tmp_path): +def test_netcdf4_entrypoint(tmp_path: Path) -> None: entrypoint = NetCDF4BackendEntrypoint() ds = create_test_data() @@ -5359,7 +5415,7 @@ def test_netcdf4_entrypoint(tmp_path): @requires_scipy -def test_scipy_entrypoint(tmp_path): +def test_scipy_entrypoint(tmp_path: Path) -> None: entrypoint = ScipyBackendEntrypoint() ds = create_test_data() @@ -5389,7 +5445,7 @@ def test_scipy_entrypoint(tmp_path): @requires_h5netcdf -def test_h5netcdf_entrypoint(tmp_path): +def test_h5netcdf_entrypoint(tmp_path: Path) -> None: entrypoint = H5netcdfBackendEntrypoint() ds = create_test_data() @@ -5427,3 +5483,51 @@ def test_write_file_from_np_str(str_type, tmpdir) -> None: txr = tdf.to_xarray() txr.to_netcdf(tmpdir.join("test.nc")) + + +@requires_zarr +@requires_netCDF4 +class TestNCZarr: + @staticmethod + def _create_nczarr(filename): + netcdfc_version = Version(nc4.getlibversion().split()[0]) + if netcdfc_version < Version("4.8.1"): + pytest.skip("requires netcdf-c>=4.8.1") + if (platform.system() == "Windows") and (netcdfc_version == Version("4.8.1")): + # Bug in netcdf-c==4.8.1 (typo: Nan instead of NaN) + # https://github.com/Unidata/netcdf-c/issues/2265 + pytest.skip("netcdf-c==4.8.1 has issues on Windows") + + ds = create_test_data() + # Drop dim3: netcdf-c does not support dtype='4.8.1 will add _ARRAY_DIMENSIONS by default + mode = "nczarr" if netcdfc_version == Version("4.8.1") else "nczarr,noxarray" + ds.to_netcdf(f"file://{filename}#mode={mode}") + return ds + + def test_open_nczarr(self): + with create_tmp_file(suffix=".zarr") as tmp: + expected = self._create_nczarr(tmp) + actual = xr.open_zarr(tmp, consolidated=False) + assert_identical(expected, actual) + + def test_overwriting_nczarr(self): + with create_tmp_file(suffix=".zarr") as tmp: + ds = self._create_nczarr(tmp) + expected = ds[["var1"]] + expected.to_zarr(tmp, mode="w") + actual = xr.open_zarr(tmp, consolidated=False) + assert_identical(expected, actual) + + @pytest.mark.parametrize("mode", ["a", "r+"]) + @pytest.mark.filterwarnings("ignore:.*non-consolidated metadata.*") + def test_raise_writing_to_nczarr(self, mode): + with create_tmp_file(suffix=".zarr") as tmp: + ds = self._create_nczarr(tmp) + with pytest.raises( + KeyError, match="missing the attribute `_ARRAY_DIMENSIONS`," + ): + ds.to_zarr(tmp, mode=mode) diff --git a/xarray/tests/test_backends_api.py b/xarray/tests/test_backends_api.py index 352ec6c10f1..0ba446818e5 100644 --- a/xarray/tests/test_backends_api.py +++ b/xarray/tests/test_backends_api.py @@ -1,9 +1,18 @@ +from numbers import Number + import numpy as np +import pytest import xarray as xr from xarray.backends.api import _get_default_engine -from . import assert_identical, requires_netCDF4, requires_scipy +from . import ( + assert_identical, + assert_no_warnings, + requires_dask, + requires_netCDF4, + requires_scipy, +) @requires_netCDF4 @@ -35,3 +44,136 @@ def open_dataset( actual = xr.open_dataset("fake_filename", engine=CustomBackend) assert_identical(expected, actual) + + +class PassThroughBackendEntrypoint(xr.backends.BackendEntrypoint): + """Access an object passed to the `open_dataset` method.""" + + def open_dataset(self, dataset, *, drop_variables=None): + """Return the first argument.""" + return dataset + + +def explicit_chunks(chunks, shape): + """Return explicit chunks, expanding any integer member to a tuple of integers.""" + # Emulate `dask.array.core.normalize_chunks` but for simpler inputs. + return tuple( + ( + (size // chunk) * (chunk,) + + ((size % chunk,) if size % chunk or size == 0 else ()) + ) + if isinstance(chunk, Number) + else chunk + for chunk, size in zip(chunks, shape) + ) + + +@requires_dask +class TestPreferredChunks: + """Test behaviors related to the backend's preferred chunks.""" + + var_name = "data" + + def create_dataset(self, shape, pref_chunks): + """Return a dataset with a variable with the given shape and preferred chunks.""" + dims = tuple(f"dim_{idx}" for idx in range(len(shape))) + return xr.Dataset( + { + self.var_name: xr.Variable( + dims, + np.empty(shape, dtype=np.dtype("V1")), + encoding={"preferred_chunks": dict(zip(dims, pref_chunks))}, + ) + } + ) + + def check_dataset(self, initial, final, expected_chunks): + assert_identical(initial, final) + assert final[self.var_name].chunks == expected_chunks + + @pytest.mark.parametrize( + "shape,pref_chunks", + [ + # Represent preferred chunking with int. + ((5,), (2,)), + # Represent preferred chunking with tuple. + ((5,), ((2, 2, 1),)), + # Represent preferred chunking with int in two dims. + ((5, 6), (4, 2)), + # Represent preferred chunking with tuple in second dim. + ((5, 6), (4, (2, 2, 2))), + ], + ) + @pytest.mark.parametrize("request_with_empty_map", [False, True]) + def test_honor_chunks(self, shape, pref_chunks, request_with_empty_map): + """Honor the backend's preferred chunks when opening a dataset.""" + initial = self.create_dataset(shape, pref_chunks) + # To keep the backend's preferred chunks, the `chunks` argument must be an + # empty mapping or map dimensions to `None`. + chunks = ( + {} + if request_with_empty_map + else dict.fromkeys(initial[self.var_name].dims, None) + ) + final = xr.open_dataset( + initial, engine=PassThroughBackendEntrypoint, chunks=chunks + ) + self.check_dataset(initial, final, explicit_chunks(pref_chunks, shape)) + + @pytest.mark.parametrize( + "shape,pref_chunks,req_chunks", + [ + # Preferred chunking is int; requested chunking is int. + ((5,), (2,), (3,)), + # Preferred chunking is int; requested chunking is tuple. + ((5,), (2,), ((2, 1, 1, 1),)), + # Preferred chunking is tuple; requested chunking is int. + ((5,), ((2, 2, 1),), (3,)), + # Preferred chunking is tuple; requested chunking is tuple. + ((5,), ((2, 2, 1),), ((2, 1, 1, 1),)), + # Split chunks along a dimension other than the first. + ((1, 5), (1, 2), (1, 3)), + ], + ) + def test_split_chunks(self, shape, pref_chunks, req_chunks): + """Warn when the requested chunks separate the backend's preferred chunks.""" + initial = self.create_dataset(shape, pref_chunks) + with pytest.warns(UserWarning): + final = xr.open_dataset( + initial, + engine=PassThroughBackendEntrypoint, + chunks=dict(zip(initial[self.var_name].dims, req_chunks)), + ) + self.check_dataset(initial, final, explicit_chunks(req_chunks, shape)) + + @pytest.mark.parametrize( + "shape,pref_chunks,req_chunks", + [ + # Keep preferred chunks using int representation. + ((5,), (2,), (2,)), + # Keep preferred chunks using tuple representation. + ((5,), (2,), ((2, 2, 1),)), + # Join chunks, leaving a final short chunk. + ((5,), (2,), (4,)), + # Join all chunks with an int larger than the dimension size. + ((5,), (2,), (6,)), + # Join one chunk using tuple representation. + ((5,), (1,), ((1, 1, 2, 1),)), + # Join one chunk using int representation. + ((5,), ((1, 1, 2, 1),), (2,)), + # Join multiple chunks using tuple representation. + ((5,), ((1, 1, 2, 1),), ((2, 3),)), + # Join chunks in multiple dimensions. + ((5, 5), (2, (1, 1, 2, 1)), (4, (2, 3))), + ], + ) + def test_join_chunks(self, shape, pref_chunks, req_chunks): + """Don't warn when the requested chunks join or keep the preferred chunks.""" + initial = self.create_dataset(shape, pref_chunks) + with assert_no_warnings(): + final = xr.open_dataset( + initial, + engine=PassThroughBackendEntrypoint, + chunks=dict(zip(initial[self.var_name].dims, req_chunks)), + ) + self.check_dataset(initial, final, explicit_chunks(req_chunks, shape)) diff --git a/xarray/tests/test_coding_times.py b/xarray/tests/test_coding_times.py index 92d27f22eb8..a5344fe4c85 100644 --- a/xarray/tests/test_coding_times.py +++ b/xarray/tests/test_coding_times.py @@ -1007,12 +1007,9 @@ def test_decode_ambiguous_time_warns(calendar) -> None: units = "days since 1-1-1" expected = num2date(dates, units, calendar=calendar, only_use_cftime_datetimes=True) - exp_warn_type = SerializationWarning if is_standard_calendar else None - - with pytest.warns(exp_warn_type) as record: - result = decode_cf_datetime(dates, units, calendar=calendar) - if is_standard_calendar: + with pytest.warns(SerializationWarning) as record: + result = decode_cf_datetime(dates, units, calendar=calendar) relevant_warnings = [ r for r in record.list @@ -1020,7 +1017,8 @@ def test_decode_ambiguous_time_warns(calendar) -> None: ] assert len(relevant_warnings) == 1 else: - assert not record + with assert_no_warnings(): + result = decode_cf_datetime(dates, units, calendar=calendar) np.testing.assert_array_equal(result, expected) @@ -1123,3 +1121,30 @@ def test_should_cftime_be_used_target_not_npable(): ValueError, match="Calendar 'noleap' is only valid with cftime." ): _should_cftime_be_used(src, "noleap", False) + + +@pytest.mark.parametrize("dtype", [np.uint8, np.uint16, np.uint32, np.uint64]) +def test_decode_cf_datetime_uint(dtype): + units = "seconds since 2018-08-22T03:23:03Z" + num_dates = dtype(50) + result = decode_cf_datetime(num_dates, units) + expected = np.asarray(np.datetime64("2018-08-22T03:23:53", "ns")) + np.testing.assert_equal(result, expected) + + +@requires_cftime +def test_decode_cf_datetime_uint64_with_cftime(): + units = "days since 1700-01-01" + num_dates = np.uint64(182621) + result = decode_cf_datetime(num_dates, units) + expected = np.asarray(np.datetime64("2200-01-01", "ns")) + np.testing.assert_equal(result, expected) + + +@requires_cftime +def test_decode_cf_datetime_uint64_with_cftime_overflow_error(): + units = "microseconds since 1700-01-01" + calendar = "360_day" + num_dates = np.uint64(1_000_000 * 86_400 * 360 * 500_000) + with pytest.raises(OverflowError): + decode_cf_datetime(num_dates, units, calendar) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 6a86738ab2f..1eaa772206e 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import functools import operator import pickle @@ -22,8 +24,15 @@ unified_dim_sizes, ) from xarray.core.pycompat import dask_version - -from . import has_dask, raise_if_dask_computes, requires_dask +from xarray.core.types import T_Xarray + +from . import ( + has_cftime, + has_dask, + raise_if_dask_computes, + requires_cftime, + requires_dask, +) def assert_identical(a, b): @@ -1928,38 +1937,199 @@ def test_where_attrs() -> None: expected = xr.DataArray([1, 0], dims="x", attrs={"attr": "x"}) assert_identical(expected, actual) + # ensure keep_attrs can handle scalar values + actual = xr.where(cond, 1, 0, keep_attrs=True) + assert actual.attrs == {} -@pytest.mark.parametrize("use_dask", [True, False]) -@pytest.mark.parametrize("use_datetime", [True, False]) -def test_polyval(use_dask, use_datetime) -> None: - if use_dask and not has_dask: - pytest.skip("requires dask") - - if use_datetime: - xcoord = xr.DataArray( - pd.date_range("2000-01-01", freq="D", periods=10), dims=("x",), name="x" - ) - x = xr.core.missing.get_clean_interp_index(xcoord, "x") - else: - x = np.arange(10) - xcoord = xr.DataArray(x, dims=("x",), name="x") - da = xr.DataArray( - np.stack((1.0 + x + 2.0 * x**2, 1.0 + 2.0 * x + 3.0 * x**2)), - dims=("d", "x"), - coords={"x": xcoord, "d": [0, 1]}, - ) - coeffs = xr.DataArray( - [[2, 1, 1], [3, 2, 1]], - dims=("d", "degree"), - coords={"d": [0, 1], "degree": [2, 1, 0]}, +@pytest.mark.parametrize( + "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] +) +@pytest.mark.parametrize( + ["x", "coeffs", "expected"], + [ + pytest.param( + xr.DataArray([1, 2, 3], dims="x"), + xr.DataArray([2, 3, 4], dims="degree", coords={"degree": [0, 1, 2]}), + xr.DataArray([9, 2 + 6 + 16, 2 + 9 + 36], dims="x"), + id="simple", + ), + pytest.param( + xr.DataArray([1, 2, 3], dims="x"), + xr.DataArray( + [[0, 1], [0, 1]], dims=("y", "degree"), coords={"degree": [0, 1]} + ), + xr.DataArray([[1, 1], [2, 2], [3, 3]], dims=("x", "y")), + id="broadcast-x", + ), + pytest.param( + xr.DataArray([1, 2, 3], dims="x"), + xr.DataArray( + [[0, 1], [1, 0], [1, 1]], + dims=("x", "degree"), + coords={"degree": [0, 1]}, + ), + xr.DataArray([1, 1, 1 + 3], dims="x"), + id="shared-dim", + ), + pytest.param( + xr.DataArray([1, 2, 3], dims="x"), + xr.DataArray([1, 0, 0], dims="degree", coords={"degree": [2, 1, 0]}), + xr.DataArray([1, 2**2, 3**2], dims="x"), + id="reordered-index", + ), + pytest.param( + xr.DataArray([1, 2, 3], dims="x"), + xr.DataArray([5], dims="degree", coords={"degree": [3]}), + xr.DataArray([5, 5 * 2**3, 5 * 3**3], dims="x"), + id="sparse-index", + ), + pytest.param( + xr.DataArray([1, 2, 3], dims="x"), + xr.Dataset( + {"a": ("degree", [0, 1]), "b": ("degree", [1, 0])}, + coords={"degree": [0, 1]}, + ), + xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [1, 1, 1])}), + id="array-dataset", + ), + pytest.param( + xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("x", [2, 3, 4])}), + xr.DataArray([1, 1], dims="degree", coords={"degree": [0, 1]}), + xr.Dataset({"a": ("x", [2, 3, 4]), "b": ("x", [3, 4, 5])}), + id="dataset-array", + ), + pytest.param( + xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("y", [2, 3, 4])}), + xr.Dataset( + {"a": ("degree", [0, 1]), "b": ("degree", [1, 1])}, + coords={"degree": [0, 1]}, + ), + xr.Dataset({"a": ("x", [1, 2, 3]), "b": ("y", [3, 4, 5])}), + id="dataset-dataset", + ), + pytest.param( + xr.DataArray(pd.date_range("1970-01-01", freq="s", periods=3), dims="x"), + xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}), + xr.DataArray( + [0, 1e9, 2e9], + dims="x", + coords={"x": pd.date_range("1970-01-01", freq="s", periods=3)}, + ), + id="datetime", + ), + pytest.param( + xr.DataArray( + np.array([1000, 2000, 3000], dtype="timedelta64[ns]"), dims="x" + ), + xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}), + xr.DataArray([1000.0, 2000.0, 3000.0], dims="x"), + id="timedelta", + ), + ], +) +def test_polyval( + use_dask: bool, + x: xr.DataArray | xr.Dataset, + coeffs: xr.DataArray | xr.Dataset, + expected: xr.DataArray | xr.Dataset, +) -> None: + if use_dask: + if not has_dask: + pytest.skip("requires dask") + coeffs = coeffs.chunk({"degree": 2}) + x = x.chunk({"x": 2}) + + with raise_if_dask_computes(): + actual = xr.polyval(coord=x, coeffs=coeffs) + + xr.testing.assert_allclose(actual, expected) + + +@requires_cftime +@pytest.mark.parametrize( + "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] +) +@pytest.mark.parametrize("date", ["1970-01-01", "0753-04-21"]) +def test_polyval_cftime(use_dask: bool, date: str) -> None: + import cftime + + x = xr.DataArray( + xr.date_range(date, freq="1S", periods=3, use_cftime=True), + dims="x", ) + coeffs = xr.DataArray([0, 1], dims="degree", coords={"degree": [0, 1]}) + if use_dask: - coeffs = coeffs.chunk({"d": 2}) + if not has_dask: + pytest.skip("requires dask") + coeffs = coeffs.chunk({"degree": 2}) + x = x.chunk({"x": 2}) + + with raise_if_dask_computes(max_computes=1): + actual = xr.polyval(coord=x, coeffs=coeffs) + + t0 = xr.date_range(date, periods=1)[0] + offset = (t0 - cftime.DatetimeGregorian(1970, 1, 1)).total_seconds() * 1e9 + expected = ( + xr.DataArray( + [0, 1e9, 2e9], + dims="x", + coords={"x": xr.date_range(date, freq="1S", periods=3, use_cftime=True)}, + ) + + offset + ) + xr.testing.assert_allclose(actual, expected) + - da_pv = xr.polyval(da.x, coeffs) +def test_polyval_degree_dim_checks() -> None: + x = xr.DataArray([1, 2, 3], dims="x") + coeffs = xr.DataArray([2, 3, 4], dims="degree", coords={"degree": [0, 1, 2]}) + with pytest.raises(ValueError): + xr.polyval(x, coeffs.drop_vars("degree")) + with pytest.raises(ValueError): + xr.polyval(x, coeffs.assign_coords(degree=coeffs.degree.astype(float))) + + +@pytest.mark.parametrize( + "use_dask", [pytest.param(False, id="nodask"), pytest.param(True, id="dask")] +) +@pytest.mark.parametrize( + "x", + [ + pytest.param(xr.DataArray([0, 1, 2], dims="x"), id="simple"), + pytest.param( + xr.DataArray(pd.date_range("1970-01-01", freq="ns", periods=3), dims="x"), + id="datetime", + ), + pytest.param( + xr.DataArray(np.array([0, 1, 2], dtype="timedelta64[ns]"), dims="x"), + id="timedelta", + ), + ], +) +@pytest.mark.parametrize( + "y", + [ + pytest.param(xr.DataArray([1, 6, 17], dims="x"), id="1D"), + pytest.param( + xr.DataArray([[1, 6, 17], [34, 57, 86]], dims=("y", "x")), id="2D" + ), + ], +) +def test_polyfit_polyval_integration( + use_dask: bool, x: xr.DataArray, y: xr.DataArray +) -> None: + y.coords["x"] = x + if use_dask: + if not has_dask: + pytest.skip("requires dask") + y = y.chunk({"x": 2}) - xr.testing.assert_allclose(da, da_pv.T) + fit = y.polyfit(dim="x", deg=2) + evaluated = xr.polyval(y.x, fit.polyfit_coefficients) + expected = y.transpose(*evaluated.dims) + xr.testing.assert_allclose(evaluated.variable, expected.variable) @pytest.mark.parametrize("use_dask", [False, True]) diff --git a/xarray/tests/test_concat.py b/xarray/tests/test_concat.py index c0b98e28b12..28973e20cd0 100644 --- a/xarray/tests/test_concat.py +++ b/xarray/tests/test_concat.py @@ -1,5 +1,5 @@ from copy import deepcopy -from typing import List +from typing import TYPE_CHECKING, Any, List import numpy as np import pandas as pd @@ -18,6 +18,9 @@ ) from .test_dataset import create_test_data +if TYPE_CHECKING: + from xarray.core.types import CombineAttrsOptions, JoinOptions + def test_concat_compat() -> None: ds1 = Dataset( @@ -239,7 +242,7 @@ def test_concat_join_kwarg(self) -> None: ds1 = Dataset({"a": (("x", "y"), [[0]])}, coords={"x": [0], "y": [0]}) ds2 = Dataset({"a": (("x", "y"), [[0]])}, coords={"x": [1], "y": [0.0001]}) - expected = {} + expected: dict[JoinOptions, Any] = {} expected["outer"] = Dataset( {"a": (("x", "y"), [[0, np.nan], [np.nan, 0]])}, {"x": [0, 1], "y": [0, 0.0001]}, @@ -456,7 +459,7 @@ def test_concat_promote_shape(self) -> None: ) assert_identical(actual, expected) - # regression GH6416 (coord dtype) + # regression GH6416 (coord dtype) and GH6434 time_data1 = np.array(["2022-01-01", "2022-02-01"], dtype="datetime64[ns]") time_data2 = np.array("2022-03-01", dtype="datetime64[ns]") time_expected = np.array( @@ -466,6 +469,7 @@ def test_concat_promote_shape(self) -> None: actual = concat(objs, "time") expected = Dataset({}, {"time": time_expected}) assert_identical(actual, expected) + assert isinstance(actual.indexes["time"], pd.DatetimeIndex) def test_concat_do_not_promote(self) -> None: # GH438 @@ -653,7 +657,7 @@ def test_concat_join_kwarg(self) -> None: {"a": (("x", "y"), [[0]])}, coords={"x": [1], "y": [0.0001]} ).to_array() - expected = {} + expected: dict[JoinOptions, Any] = {} expected["outer"] = Dataset( {"a": (("x", "y"), [[0, np.nan], [np.nan, 0]])}, {"x": [0, 1], "y": [0, 0.0001]}, @@ -685,7 +689,7 @@ def test_concat_combine_attrs_kwarg(self) -> None: da1 = DataArray([0], coords=[("x", [0])], attrs={"b": 42}) da2 = DataArray([0], coords=[("x", [1])], attrs={"b": 42, "c": 43}) - expected = {} + expected: dict[CombineAttrsOptions, Any] = {} expected["drop"] = DataArray([0, 0], coords=[("x", [0, 1])]) expected["no_conflicts"] = DataArray( [0, 0], coords=[("x", [0, 1])], attrs={"b": 42, "c": 43} diff --git a/xarray/tests/test_conventions.py b/xarray/tests/test_conventions.py index 83e560e7208..b8b9d19e238 100644 --- a/xarray/tests/test_conventions.py +++ b/xarray/tests/test_conventions.py @@ -9,6 +9,7 @@ Dataset, SerializationWarning, Variable, + cftime_range, coding, conventions, open_dataset, @@ -128,6 +129,25 @@ def test_multidimensional_coordinates(self) -> None: # Should not have any global coordinates. assert "coordinates" not in attrs + def test_var_with_coord_attr(self) -> None: + # regression test for GH6310 + # don't overwrite user-defined "coordinates" attributes + orig = Dataset( + {"values": ("time", np.zeros(2), {"coordinates": "time lon lat"})}, + coords={ + "time": ("time", np.zeros(2)), + "lat": ("time", np.zeros(2)), + "lon": ("time", np.zeros(2)), + }, + ) + # Encode the coordinates, as they would be in a netCDF output file. + enc, attrs = conventions.encode_dataset_coordinates(orig) + # Make sure we have the right coordinates for each variable. + values_coords = enc["values"].attrs.get("coordinates", "") + assert set(values_coords.split()) == {"time", "lat", "lon"} + # Should not have any global coordinates. + assert "coordinates" not in attrs + def test_do_not_overwrite_user_coordinates(self) -> None: orig = Dataset( coords={"x": [0, 1, 2], "y": ("x", [5, 6, 7]), "z": ("x", [8, 9, 10])}, @@ -423,3 +443,25 @@ def test_decode_cf_variable_with_array_units(self) -> None: v = Variable(["t"], [1, 2, 3], {"units": np.array(["foobar"], dtype=object)}) v_decoded = conventions.decode_cf_variable("test2", v) assert_identical(v, v_decoded) + + +def test_decode_cf_variable_timedelta64(): + variable = Variable(["time"], pd.timedelta_range("1D", periods=2)) + decoded = conventions.decode_cf_variable("time", variable) + assert decoded.encoding == {} + assert_identical(decoded, variable) + + +def test_decode_cf_variable_datetime64(): + variable = Variable(["time"], pd.date_range("2000", periods=2)) + decoded = conventions.decode_cf_variable("time", variable) + assert decoded.encoding == {} + assert_identical(decoded, variable) + + +@requires_cftime +def test_decode_cf_variable_cftime(): + variable = Variable(["time"], cftime_range("2000", periods=2)) + decoded = conventions.decode_cf_variable("time", variable) + assert decoded.encoding == {} + assert_identical(decoded, variable) diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 872c0c6f1db..df69e8d9d6e 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -10,7 +10,6 @@ from packaging.version import Version import xarray as xr -import xarray.ufuncs as xu from xarray import DataArray, Dataset, Variable from xarray.core import duck_array_ops from xarray.core.pycompat import dask_version @@ -265,18 +264,16 @@ def test_missing_methods(self): except NotImplementedError as err: assert "dask" in str(err) - @pytest.mark.filterwarnings("ignore::FutureWarning") def test_univariate_ufunc(self): u = self.eager_var v = self.lazy_var - self.assertLazyAndAllClose(np.sin(u), xu.sin(v)) + self.assertLazyAndAllClose(np.sin(u), np.sin(v)) - @pytest.mark.filterwarnings("ignore::FutureWarning") def test_bivariate_ufunc(self): u = self.eager_var v = self.lazy_var - self.assertLazyAndAllClose(np.maximum(u, 0), xu.maximum(v, 0)) - self.assertLazyAndAllClose(np.maximum(u, 0), xu.maximum(0, v)) + self.assertLazyAndAllClose(np.maximum(u, 0), np.maximum(v, 0)) + self.assertLazyAndAllClose(np.maximum(u, 0), np.maximum(0, v)) def test_compute(self): u = self.eager_var @@ -605,11 +602,10 @@ def duplicate_and_merge(array): actual = duplicate_and_merge(self.lazy_array) self.assertLazyAndEqual(expected, actual) - @pytest.mark.filterwarnings("ignore::FutureWarning") def test_ufuncs(self): u = self.eager_array v = self.lazy_array - self.assertLazyAndAllClose(np.sin(u), xu.sin(v)) + self.assertLazyAndAllClose(np.sin(u), np.sin(v)) def test_where_dispatching(self): a = np.arange(10) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 65efb3a732c..e488f5afad9 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -3,6 +3,7 @@ import warnings from copy import deepcopy from textwrap import dedent +from typing import Any, Final, Hashable, cast import numpy as np import pandas as pd @@ -25,6 +26,7 @@ from xarray.core import dtypes from xarray.core.common import full_like from xarray.core.indexes import Index, PandasIndex, filter_indexes_from_coords +from xarray.core.types import QueryEngineOptions, QueryParserOptions from xarray.core.utils import is_scalar from xarray.tests import ( ReturnItem, @@ -69,7 +71,7 @@ def setup(self): ) self.mda = DataArray([0, 1, 2, 3], coords={"x": self.mindex}, dims="x") - def test_repr(self): + def test_repr(self) -> None: v = Variable(["time", "x"], [[1, 2, 3], [4, 5, 6]], {"foo": "bar"}) coords = {"x": np.arange(3, dtype=np.int64), "other": np.int64(0)} data_array = DataArray(v, coords, name="my_variable") @@ -87,7 +89,7 @@ def test_repr(self): ) assert expected == repr(data_array) - def test_repr_multiindex(self): + def test_repr_multiindex(self) -> None: expected = dedent( """\ @@ -99,7 +101,7 @@ def test_repr_multiindex(self): ) assert expected == repr(self.mda) - def test_repr_multiindex_long(self): + def test_repr_multiindex_long(self) -> None: mindex_long = pd.MultiIndex.from_product( [["a", "b", "c", "d"], [1, 2, 3, 4, 5, 6, 7, 8]], names=("level_1", "level_2"), @@ -117,7 +119,7 @@ def test_repr_multiindex_long(self): ) assert expected == repr(mda_long) - def test_properties(self): + def test_properties(self) -> None: assert_equal(self.dv.variable, self.v) assert_array_equal(self.dv.values, self.v.values) for attr in ["dims", "dtype", "shape", "size", "nbytes", "ndim", "attrs"]: @@ -135,7 +137,7 @@ def test_properties(self): with pytest.raises(AttributeError): self.dv.variable = self.v - def test_data_property(self): + def test_data_property(self) -> None: array = DataArray(np.zeros((3, 4))) actual = array.copy() actual.values = np.ones((3, 4)) @@ -144,7 +146,7 @@ def test_data_property(self): assert_array_equal(2 * np.ones((3, 4)), actual.data) assert_array_equal(actual.data, actual.values) - def test_indexes(self): + def test_indexes(self) -> None: array = DataArray(np.zeros((2, 3)), [("x", [0, 1]), ("y", ["a", "b", "c"])]) expected_indexes = {"x": pd.Index([0, 1]), "y": pd.Index(["a", "b", "c"])} expected_xindexes = { @@ -158,21 +160,21 @@ def test_indexes(self): assert array.xindexes[k].equals(expected_xindexes[k]) assert array.indexes[k].equals(expected_indexes[k]) - def test_get_index(self): + def test_get_index(self) -> None: array = DataArray(np.zeros((2, 3)), coords={"x": ["a", "b"]}, dims=["x", "y"]) assert array.get_index("x").equals(pd.Index(["a", "b"])) assert array.get_index("y").equals(pd.Index([0, 1, 2])) with pytest.raises(KeyError): array.get_index("z") - def test_get_index_size_zero(self): + def test_get_index_size_zero(self) -> None: array = DataArray(np.zeros((0,)), dims=["x"]) actual = array.get_index("x") expected = pd.Index([], dtype=np.int64) assert actual.equals(expected) assert actual.dtype == expected.dtype - def test_struct_array_dims(self): + def test_struct_array_dims(self) -> None: """ This test checks subtraction of two DataArrays for the case when dimension is a structured array. @@ -230,7 +232,7 @@ def test_struct_array_dims(self): assert_identical(actual, expected) - def test_name(self): + def test_name(self) -> None: arr = self.dv assert arr.name == "foo" @@ -244,33 +246,33 @@ def test_name(self): expected = DataArray([3], [("x", [3])], name="y") assert_identical(actual, expected) - def test_dims(self): + def test_dims(self) -> None: arr = self.dv assert arr.dims == ("x", "y") with pytest.raises(AttributeError, match=r"you cannot assign"): arr.dims = ("w", "z") - def test_sizes(self): + def test_sizes(self) -> None: array = DataArray(np.zeros((3, 4)), dims=["x", "y"]) assert array.sizes == {"x": 3, "y": 4} assert tuple(array.sizes) == array.dims with pytest.raises(TypeError): - array.sizes["foo"] = 5 + array.sizes["foo"] = 5 # type: ignore - def test_encoding(self): + def test_encoding(self) -> None: expected = {"foo": "bar"} self.dv.encoding["foo"] = "bar" assert expected == self.dv.encoding - expected = {"baz": 0} - self.dv.encoding = expected + expected2 = {"baz": 0} + self.dv.encoding = expected2 + assert expected2 is not self.dv.encoding - assert expected is not self.dv.encoding - - def test_constructor(self): + def test_constructor(self) -> None: data = np.random.random((2, 3)) + # w/o coords, w/o dims actual = DataArray(data) expected = Dataset({None: (["dim_0", "dim_1"], data)})[None] assert_identical(expected, actual) @@ -285,6 +287,7 @@ def test_constructor(self): )[None] assert_identical(expected, actual) + # pd.Index coords, w/o dims actual = DataArray( data, [pd.Index(["a", "b"], name="x"), pd.Index([-1, -2, -3], name="y")] ) @@ -293,54 +296,66 @@ def test_constructor(self): )[None] assert_identical(expected, actual) - coords = [["a", "b"], [-1, -2, -3]] - actual = DataArray(data, coords, ["x", "y"]) + # list coords, w dims + coords1 = [["a", "b"], [-1, -2, -3]] + actual = DataArray(data, coords1, ["x", "y"]) assert_identical(expected, actual) - coords = [pd.Index(["a", "b"], name="A"), pd.Index([-1, -2, -3], name="B")] - actual = DataArray(data, coords, ["x", "y"]) + # pd.Index coords, w dims + coords2 = [pd.Index(["a", "b"], name="A"), pd.Index([-1, -2, -3], name="B")] + actual = DataArray(data, coords2, ["x", "y"]) assert_identical(expected, actual) - coords = {"x": ["a", "b"], "y": [-1, -2, -3]} - actual = DataArray(data, coords, ["x", "y"]) + # dict coords, w dims + coords3 = {"x": ["a", "b"], "y": [-1, -2, -3]} + actual = DataArray(data, coords3, ["x", "y"]) assert_identical(expected, actual) - actual = DataArray(data, coords) + # dict coords, w/o dims + actual = DataArray(data, coords3) assert_identical(expected, actual) - coords = [("x", ["a", "b"]), ("y", [-1, -2, -3])] - actual = DataArray(data, coords) + # tuple[dim, list] coords, w/o dims + coords4 = [("x", ["a", "b"]), ("y", [-1, -2, -3])] + actual = DataArray(data, coords4) assert_identical(expected, actual) + # partial dict coords, w dims expected = Dataset({None: (["x", "y"], data), "x": ("x", ["a", "b"])})[None] actual = DataArray(data, {"x": ["a", "b"]}, ["x", "y"]) assert_identical(expected, actual) + # w/o coords, w dims actual = DataArray(data, dims=["x", "y"]) expected = Dataset({None: (["x", "y"], data)})[None] assert_identical(expected, actual) + # w/o coords, w dims, w name actual = DataArray(data, dims=["x", "y"], name="foo") expected = Dataset({"foo": (["x", "y"], data)})["foo"] assert_identical(expected, actual) + # w/o coords, w/o dims, w name actual = DataArray(data, name="foo") expected = Dataset({"foo": (["dim_0", "dim_1"], data)})["foo"] assert_identical(expected, actual) + # w/o coords, w dims, w attrs actual = DataArray(data, dims=["x", "y"], attrs={"bar": 2}) expected = Dataset({None: (["x", "y"], data, {"bar": 2})})[None] assert_identical(expected, actual) + # w/o coords, w dims (ds has attrs) actual = DataArray(data, dims=["x", "y"]) expected = Dataset({None: (["x", "y"], data, {}, {"bar": 2})})[None] assert_identical(expected, actual) + # data is list, w coords actual = DataArray([1, 2, 3], coords={"x": [0, 1, 2]}) expected = DataArray([1, 2, 3], coords=[("x", [0, 1, 2])]) assert_identical(expected, actual) - def test_constructor_invalid(self): + def test_constructor_invalid(self) -> None: data = np.random.randn(3, 2) with pytest.raises(ValueError, match=r"coords is not dict-like"): @@ -367,7 +382,7 @@ def test_constructor_invalid(self): with pytest.raises(ValueError, match=r"matching the dimension size"): DataArray(data, coords={"x": 0}, dims=["x", "y"]) - def test_constructor_from_self_described(self): + def test_constructor_from_self_described(self) -> None: data = [[-0.1, 21], [0, 2]] expected = DataArray( data, @@ -413,7 +428,7 @@ def test_constructor_from_self_described(self): assert_identical(expected, actual) @requires_dask - def test_constructor_from_self_described_chunked(self): + def test_constructor_from_self_described_chunked(self) -> None: expected = DataArray( [[-0.1, 21], [0, 2]], coords={"x": ["a", "b"], "y": [-1, -2]}, @@ -425,13 +440,13 @@ def test_constructor_from_self_described_chunked(self): assert_identical(expected, actual) assert_chunks_equal(expected, actual) - def test_constructor_from_0d(self): + def test_constructor_from_0d(self) -> None: expected = Dataset({None: ([], 0)})[None] actual = DataArray(0) assert_identical(expected, actual) @requires_dask - def test_constructor_dask_coords(self): + def test_constructor_dask_coords(self) -> None: # regression test for GH1684 import dask.array as da @@ -443,7 +458,7 @@ def test_constructor_dask_coords(self): expected = DataArray(data, coords={"x": ecoord, "y": ecoord}, dims=["x", "y"]) assert_equal(actual, expected) - def test_equals_and_identical(self): + def test_equals_and_identical(self) -> None: orig = DataArray(np.arange(5.0), {"a": 42}, dims="x") expected = orig @@ -488,13 +503,13 @@ def test_equals_and_identical(self): assert not expected.equals(actual) assert not expected.identical(actual) - def test_equals_failures(self): + def test_equals_failures(self) -> None: orig = DataArray(np.arange(5.0), {"a": 42}, dims="x") - assert not orig.equals(np.arange(5)) - assert not orig.identical(123) - assert not orig.broadcast_equals({1: 2}) + assert not orig.equals(np.arange(5)) # type: ignore + assert not orig.identical(123) # type: ignore + assert not orig.broadcast_equals({1: 2}) # type: ignore - def test_broadcast_equals(self): + def test_broadcast_equals(self) -> None: a = DataArray([0, 0], {"y": 0}, dims="x") b = DataArray([0, 0], {"y": ("x", [0, 0])}, dims="x") assert a.broadcast_equals(b) @@ -506,7 +521,7 @@ def test_broadcast_equals(self): assert not a.broadcast_equals(c) assert not c.broadcast_equals(a) - def test_getitem(self): + def test_getitem(self) -> None: # strings pull out dataarrays assert_identical(self.dv, self.ds["foo"]) x = self.dv["x"] @@ -543,12 +558,12 @@ def test_getitem(self): ]: assert_array_equal(self.v[i], self.dv[i]) - def test_getitem_dict(self): + def test_getitem_dict(self) -> None: actual = self.dv[{"x": slice(3), "y": 0}] expected = self.dv.isel(x=slice(3), y=0) assert_identical(expected, actual) - def test_getitem_coords(self): + def test_getitem_coords(self) -> None: orig = DataArray( [[10], [20]], { @@ -604,7 +619,7 @@ def test_getitem_coords(self): ) assert_identical(expected, actual) - def test_getitem_dataarray(self): + def test_getitem_dataarray(self) -> None: # It should not conflict da = DataArray(np.arange(12).reshape((3, 4)), dims=["x", "y"]) ind = DataArray([[0, 1], [0, 1]], dims=["x", "z"]) @@ -628,7 +643,7 @@ def test_getitem_dataarray(self): assert_equal(da[ind], da[[0, 1]]) assert_equal(da[ind], da[ind.values]) - def test_getitem_empty_index(self): + def test_getitem_empty_index(self) -> None: da = DataArray(np.arange(12).reshape((3, 4)), dims=["x", "y"]) assert_identical(da[{"x": []}], DataArray(np.zeros((0, 4)), dims=["x", "y"])) assert_identical( @@ -636,7 +651,7 @@ def test_getitem_empty_index(self): ) assert_identical(da[[]], DataArray(np.zeros((0, 4)), dims=["x", "y"])) - def test_setitem(self): + def test_setitem(self) -> None: # basic indexing should work as numpy's indexing tuples = [ (0, 0), @@ -663,7 +678,7 @@ def test_setitem(self): expected[t] = 1 assert_array_equal(orig.values, expected) - def test_setitem_fancy(self): + def test_setitem_fancy(self) -> None: # vectorized indexing da = DataArray(np.ones((3, 2)), dims=["x", "y"]) ind = Variable(["a"], [0, 1]) @@ -693,7 +708,7 @@ def test_setitem_fancy(self): expected = DataArray([[0, 0], [0, 0], [1, 1]], dims=["x", "y"]) assert_identical(expected, da) - def test_setitem_dataarray(self): + def test_setitem_dataarray(self) -> None: def get_data(): return DataArray( np.ones((4, 3, 2)), @@ -764,18 +779,18 @@ def get_data(): ) da[dict(x=ind)] = value # should not raise - def test_contains(self): + def test_contains(self) -> None: data_array = DataArray([1, 2]) assert 1 in data_array assert 3 not in data_array - def test_pickle(self): + def test_pickle(self) -> None: data = DataArray(np.random.random((3, 3)), dims=("id", "time")) roundtripped = pickle.loads(pickle.dumps(data)) assert_identical(data, roundtripped) @requires_dask - def test_chunk(self): + def test_chunk(self) -> None: unblocked = DataArray(np.ones((3, 4))) assert unblocked.chunks is None @@ -804,7 +819,12 @@ def test_chunk(self): assert isinstance(blocked.data, da.Array) assert "testname_" in blocked.data.name - def test_isel(self): + # test kwargs form of chunks + blocked = unblocked.chunk(dim_0=3, dim_1=3) + assert blocked.chunks == ((3,), (3, 1)) + assert blocked.data.name != first_dask_name + + def test_isel(self) -> None: assert_identical(self.dv[0], self.dv.isel(x=0)) assert_identical(self.dv, self.dv.isel(x=slice(None))) assert_identical(self.dv[:3], self.dv.isel(x=slice(3))) @@ -823,7 +843,7 @@ def test_isel(self): self.dv.isel(not_a_dim=0, missing_dims="warn") assert_identical(self.dv, self.dv.isel(not_a_dim=0, missing_dims="ignore")) - def test_isel_types(self): + def test_isel_types(self) -> None: # regression test for #1405 da = DataArray([1, 2, 3], dims="x") # uint64 @@ -840,7 +860,7 @@ def test_isel_types(self): ) @pytest.mark.filterwarnings("ignore::DeprecationWarning") - def test_isel_fancy(self): + def test_isel_fancy(self) -> None: shape = (10, 7, 6) np_array = np.random.random(shape) da = DataArray( @@ -871,9 +891,9 @@ def test_isel_fancy(self): da.isel(time=(("points",), [1, 2])) y = [-1, 0] x = [-2, 2] - expected = da.values[:, y, x] - actual = da.isel(x=(("points",), x), y=(("points",), y)).values - np.testing.assert_equal(actual, expected) + expected2 = da.values[:, y, x] + actual2 = da.isel(x=(("points",), x), y=(("points",), y)).values + np.testing.assert_equal(actual2, expected2) # test that the order of the indexers doesn't matter assert_identical( @@ -891,10 +911,10 @@ def test_isel_fancy(self): stations["dim1s"] = (("station",), [1, 2, 3]) stations["dim2s"] = (("station",), [4, 5, 1]) - actual = da.isel(x=stations["dim1s"], y=stations["dim2s"]) - assert "station" in actual.coords - assert "station" in actual.dims - assert_identical(actual["station"], stations["station"]) + actual3 = da.isel(x=stations["dim1s"], y=stations["dim2s"]) + assert "station" in actual3.coords + assert "station" in actual3.dims + assert_identical(actual3["station"], stations["station"]) with pytest.raises(ValueError, match=r"conflicting values/indexes on "): da.isel( @@ -909,19 +929,19 @@ def test_isel_fancy(self): stations["dim1s"] = (("a", "b"), [[1, 2], [2, 3], [3, 4]]) stations["dim2s"] = (("a",), [4, 5, 1]) - actual = da.isel(x=stations["dim1s"], y=stations["dim2s"]) - assert "a" in actual.coords - assert "a" in actual.dims - assert "b" in actual.coords - assert "b" in actual.dims - assert_identical(actual["a"], stations["a"]) - assert_identical(actual["b"], stations["b"]) - expected = da.variable[ + actual4 = da.isel(x=stations["dim1s"], y=stations["dim2s"]) + assert "a" in actual4.coords + assert "a" in actual4.dims + assert "b" in actual4.coords + assert "b" in actual4.dims + assert_identical(actual4["a"], stations["a"]) + assert_identical(actual4["b"], stations["b"]) + expected4 = da.variable[ :, stations["dim2s"].variable, stations["dim1s"].variable ] - assert_array_equal(actual, expected) + assert_array_equal(actual4, expected4) - def test_sel(self): + def test_sel(self) -> None: self.ds["x"] = ("x", np.array(list("abcdefghij"))) da = self.ds["foo"] assert_identical(da, da.sel(x=slice(None))) @@ -934,7 +954,7 @@ def test_sel(self): assert_identical(da[1], da.sel(x=b)) assert_identical(da[[1]], da.sel(x=slice(b, b))) - def test_sel_dataarray(self): + def test_sel_dataarray(self) -> None: # indexing with DataArray self.ds["x"] = ("x", np.array(list("abcdefghij"))) da = self.ds["foo"] @@ -959,12 +979,12 @@ def test_sel_dataarray(self): assert "new_dim" in actual.coords assert_equal(actual["new_dim"].drop_vars("x"), ind["new_dim"]) - def test_sel_invalid_slice(self): + def test_sel_invalid_slice(self) -> None: array = DataArray(np.arange(10), [("x", np.arange(10))]) with pytest.raises(ValueError, match=r"cannot use non-scalar arrays"): array.sel(x=slice(array.x)) - def test_sel_dataarray_datetime_slice(self): + def test_sel_dataarray_datetime_slice(self) -> None: # regression test for GH1240 times = pd.date_range("2000-01-01", freq="D", periods=365) array = DataArray(np.arange(365), [("time", times)]) @@ -975,7 +995,7 @@ def test_sel_dataarray_datetime_slice(self): result = array.sel(delta=slice(array.delta[0], array.delta[-1])) assert_equal(result, array) - def test_sel_float(self): + def test_sel_float(self) -> None: data_values = np.arange(4) # case coords are float32 and label is list of floats @@ -1002,7 +1022,7 @@ def test_sel_float(self): assert_equal(expected_scalar, actual_scalar) assert_equal(expected_16, actual_16) - def test_sel_float_multiindex(self): + def test_sel_float_multiindex(self) -> None: # regression test https://github.com/pydata/xarray/issues/5691 # test multi-index created from coordinates, one with dtype=float32 lvl1 = ["a", "a", "b", "b"] @@ -1017,14 +1037,14 @@ def test_sel_float_multiindex(self): assert_equal(actual, expected) - def test_sel_no_index(self): + def test_sel_no_index(self) -> None: array = DataArray(np.arange(10), dims="x") assert_identical(array[0], array.sel(x=0)) assert_identical(array[:5], array.sel(x=slice(5))) assert_identical(array[[0, -1]], array.sel(x=[0, -1])) assert_identical(array[array < 5], array.sel(x=(array < 5))) - def test_sel_method(self): + def test_sel_method(self) -> None: data = DataArray(np.random.randn(3, 4), [("x", [0, 1, 2]), ("y", list("abcd"))]) expected = data.sel(y=["a", "b"]) @@ -1035,7 +1055,7 @@ def test_sel_method(self): actual = data.sel(x=[0.9, 1.9], method="backfill", tolerance=1) assert_identical(expected, actual) - def test_sel_drop(self): + def test_sel_drop(self) -> None: data = DataArray([1, 2, 3], [("x", [0, 1, 2])]) expected = DataArray(1) selected = data.sel(x=0, drop=True) @@ -1050,7 +1070,7 @@ def test_sel_drop(self): selected = data.sel(x=0, drop=True) assert_identical(expected, selected) - def test_isel_drop(self): + def test_isel_drop(self) -> None: data = DataArray([1, 2, 3], [("x", [0, 1, 2])]) expected = DataArray(1) selected = data.isel(x=0, drop=True) @@ -1060,7 +1080,7 @@ def test_isel_drop(self): selected = data.isel(x=0, drop=False) assert_identical(expected, selected) - def test_head(self): + def test_head(self) -> None: assert_equal(self.dv.isel(x=slice(5)), self.dv.head(x=5)) assert_equal(self.dv.isel(x=slice(0)), self.dv.head(x=0)) assert_equal( @@ -1076,7 +1096,7 @@ def test_head(self): with pytest.raises(ValueError, match=r"expected positive int"): self.dv.head(-3) - def test_tail(self): + def test_tail(self) -> None: assert_equal(self.dv.isel(x=slice(-5, None)), self.dv.tail(x=5)) assert_equal(self.dv.isel(x=slice(0)), self.dv.tail(x=0)) assert_equal( @@ -1093,7 +1113,7 @@ def test_tail(self): with pytest.raises(ValueError, match=r"expected positive int"): self.dv.tail(-3) - def test_thin(self): + def test_thin(self) -> None: assert_equal(self.dv.isel(x=slice(None, None, 5)), self.dv.thin(x=5)) assert_equal( self.dv.isel({dim: slice(None, None, 6) for dim in self.dv.dims}), @@ -1108,10 +1128,11 @@ def test_thin(self): with pytest.raises(ValueError, match=r"cannot be zero"): self.dv.thin(time=0) - def test_loc(self): + def test_loc(self) -> None: self.ds["x"] = ("x", np.array(list("abcdefghij"))) da = self.ds["foo"] - assert_identical(da[:3], da.loc[:"c"]) + # typing issue: see https://github.com/python/mypy/issues/2410 + assert_identical(da[:3], da.loc[:"c"]) # type: ignore[misc] assert_identical(da[1], da.loc["b"]) assert_identical(da[1], da.loc[{"x": "b"}]) assert_identical(da[1], da.loc["b", ...]) @@ -1119,17 +1140,18 @@ def test_loc(self): assert_identical(da[:3, :4], da.loc[["a", "b", "c"], np.arange(4)]) assert_identical(da[:, :4], da.loc[:, self.ds["y"] < 4]) - def test_loc_datetime64_value(self): + def test_loc_datetime64_value(self) -> None: # regression test for https://github.com/pydata/xarray/issues/4283 t = np.array(["2017-09-05T12", "2017-09-05T15"], dtype="datetime64[ns]") array = DataArray(np.ones(t.shape), dims=("time",), coords=(t,)) assert_identical(array.loc[{"time": t[0]}], array[0]) - def test_loc_assign(self): + def test_loc_assign(self) -> None: self.ds["x"] = ("x", np.array(list("abcdefghij"))) da = self.ds["foo"] # assignment - da.loc["a":"j"] = 0 + # typing issue: see https://github.com/python/mypy/issues/2410 + da.loc["a":"j"] = 0 # type: ignore[misc] assert np.all(da.values == 0) da.loc[{"x": slice("a", "j")}] = 2 assert np.all(da.values == 2) @@ -1148,7 +1170,7 @@ def test_loc_assign(self): assert np.all(da.values[0] == np.zeros(4)) assert da.values[1, 0] != 0 - def test_loc_assign_dataarray(self): + def test_loc_assign_dataarray(self) -> None: def get_data(): return DataArray( np.ones((4, 3, 2)), @@ -1194,12 +1216,12 @@ def get_data(): assert_identical(da["x"], get_data()["x"]) assert_identical(da["non-dim"], get_data()["non-dim"]) - def test_loc_single_boolean(self): + def test_loc_single_boolean(self) -> None: data = DataArray([0, 1], coords=[[True, False]]) assert data.loc[True] == 0 assert data.loc[False] == 1 - def test_loc_dim_name_collision_with_sel_params(self): + def test_loc_dim_name_collision_with_sel_params(self) -> None: da = xr.DataArray( [[0, 0], [1, 1]], dims=["dim1", "method"], @@ -1209,13 +1231,15 @@ def test_loc_dim_name_collision_with_sel_params(self): da.loc[dict(dim1=["x", "y"], method=["a"])], [[0], [1]] ) - def test_selection_multiindex(self): + def test_selection_multiindex(self) -> None: mindex = pd.MultiIndex.from_product( [["a", "b"], [1, 2], [-1, -2]], names=("one", "two", "three") ) mdata = DataArray(range(8), [("x", mindex)]) - def test_sel(lab_indexer, pos_indexer, replaced_idx=False, renamed_dim=None): + def test_sel( + lab_indexer, pos_indexer, replaced_idx=False, renamed_dim=None + ) -> None: da = mdata.sel(x=lab_indexer) expected_da = mdata.isel(x=pos_indexer) if not replaced_idx: @@ -1247,7 +1271,7 @@ def test_sel(lab_indexer, pos_indexer, replaced_idx=False, renamed_dim=None): assert_identical(mdata.sel(x={"one": "a", "two": 1}), mdata.sel(one="a", two=1)) - def test_selection_multiindex_remove_unused(self): + def test_selection_multiindex_remove_unused(self) -> None: # GH2619. For MultiIndex, we need to call remove_unused. ds = xr.DataArray( np.arange(40).reshape(8, 5), @@ -1264,7 +1288,7 @@ def test_selection_multiindex_remove_unused(self): expected = expected.set_index(xy=["x", "y"]).unstack() assert_identical(expected, actual) - def test_selection_multiindex_from_level(self): + def test_selection_multiindex_from_level(self) -> None: # GH: 3512 da = DataArray([0, 1], dims=["x"], coords={"x": [0, 1], "y": "a"}) db = DataArray([2, 3], dims=["x"], coords={"x": [0, 1], "y": "b"}) @@ -1274,20 +1298,20 @@ def test_selection_multiindex_from_level(self): expected = data.isel(xy=[0, 1]).unstack("xy").squeeze("y") assert_equal(actual, expected) - def test_virtual_default_coords(self): + def test_virtual_default_coords(self) -> None: array = DataArray(np.zeros((5,)), dims="x") expected = DataArray(range(5), dims="x", name="x") assert_identical(expected, array["x"]) assert_identical(expected, array.coords["x"]) - def test_virtual_time_components(self): + def test_virtual_time_components(self) -> None: dates = pd.date_range("2000-01-01", periods=10) da = DataArray(np.arange(1, 11), [("time", dates)]) assert_array_equal(da["time.dayofyear"], da.values) assert_array_equal(da.coords["time.dayofyear"], da.values) - def test_coords(self): + def test_coords(self) -> None: # use int64 to ensure repr() consistency on windows coords = [ IndexVariable("x", np.array([-1, -2], "int64")), @@ -1311,14 +1335,14 @@ def test_coords(self): with pytest.raises(KeyError): da.coords["foo"] - expected = dedent( + expected_repr = dedent( """\ Coordinates: * x (x) int64 -1 -2 * y (y) int64 0 1 2""" ) actual = repr(da.coords) - assert expected == actual + assert expected_repr == actual del da.coords["x"] da._indexes = filter_indexes_from_coords(da.xindexes, set(da.coords)) @@ -1331,7 +1355,7 @@ def test_coords(self): self.mda["level_1"] = ("x", np.arange(4)) self.mda.coords["level_1"] = ("x", np.arange(4)) - def test_coords_to_index(self): + def test_coords_to_index(self) -> None: da = DataArray(np.zeros((2, 3)), [("x", [1, 2]), ("y", list("abc"))]) with pytest.raises(ValueError, match=r"no valid index"): @@ -1356,7 +1380,7 @@ def test_coords_to_index(self): with pytest.raises(ValueError, match=r"ordered_dims must match"): da.coords.to_index(["x"]) - def test_coord_coords(self): + def test_coord_coords(self) -> None: orig = DataArray( [10, 20], {"x": [1, 2], "x2": ("x", ["a", "b"]), "z": 4}, dims="x" ) @@ -1376,7 +1400,7 @@ def test_coord_coords(self): ) assert_identical(expected, actual) - def test_reset_coords(self): + def test_reset_coords(self) -> None: data = DataArray( np.zeros((3, 4)), {"bar": ("x", ["a", "b", "c"]), "baz": ("y", range(4)), "y": range(4)}, @@ -1384,8 +1408,8 @@ def test_reset_coords(self): name="foo", ) - actual = data.reset_coords() - expected = Dataset( + actual1 = data.reset_coords() + expected1 = Dataset( { "foo": (["x", "y"], np.zeros((3, 4))), "bar": ("x", ["a", "b", "c"]), @@ -1393,39 +1417,38 @@ def test_reset_coords(self): "y": range(4), } ) - assert_identical(actual, expected) + assert_identical(actual1, expected1) - actual = data.reset_coords(["bar", "baz"]) - assert_identical(actual, expected) + actual2 = data.reset_coords(["bar", "baz"]) + assert_identical(actual2, expected1) - actual = data.reset_coords("bar") - expected = Dataset( + actual3 = data.reset_coords("bar") + expected3 = Dataset( {"foo": (["x", "y"], np.zeros((3, 4))), "bar": ("x", ["a", "b", "c"])}, {"baz": ("y", range(4)), "y": range(4)}, ) - assert_identical(actual, expected) + assert_identical(actual3, expected3) - actual = data.reset_coords(["bar"]) - assert_identical(actual, expected) + actual4 = data.reset_coords(["bar"]) + assert_identical(actual4, expected3) - actual = data.reset_coords(drop=True) - expected = DataArray( + actual5 = data.reset_coords(drop=True) + expected5 = DataArray( np.zeros((3, 4)), coords={"y": range(4)}, dims=["x", "y"], name="foo" ) - assert_identical(actual, expected) + assert_identical(actual5, expected5) - actual = data.copy() - actual = actual.reset_coords(drop=True) - assert_identical(actual, expected) + actual6 = data.copy().reset_coords(drop=True) + assert_identical(actual6, expected5) - actual = data.reset_coords("bar", drop=True) - expected = DataArray( + actual7 = data.reset_coords("bar", drop=True) + expected7 = DataArray( np.zeros((3, 4)), {"baz": ("y", range(4)), "y": range(4)}, dims=["x", "y"], name="foo", ) - assert_identical(actual, expected) + assert_identical(actual7, expected7) with pytest.raises(ValueError, match=r"cannot be found"): data.reset_coords("foo", drop=True) @@ -1440,7 +1463,7 @@ def test_reset_coords(self): with pytest.raises(ValueError, match=r"cannot remove index"): data.reset_coords("lvl1") - def test_assign_coords(self): + def test_assign_coords(self) -> None: array = DataArray(10) actual = array.assign_coords(c=42) expected = DataArray(10, {"c": 42}) @@ -1460,7 +1483,7 @@ def test_assign_coords(self): with pytest.raises(ValueError): da.coords["x"] = ("y", [1, 2, 3]) # no new dimension to a DataArray - def test_coords_alignment(self): + def test_coords_alignment(self) -> None: lhs = DataArray([1, 2, 3], [("x", [0, 1, 2])]) rhs = DataArray([2, 3, 4], [("x", [1, 2, 3])]) lhs.coords["rhs"] = rhs @@ -1470,18 +1493,18 @@ def test_coords_alignment(self): ) assert_identical(lhs, expected) - def test_set_coords_update_index(self): + def test_set_coords_update_index(self) -> None: actual = DataArray([1, 2, 3], [("x", [1, 2, 3])]) actual.coords["x"] = ["a", "b", "c"] assert actual.xindexes["x"].to_pandas_index().equals(pd.Index(["a", "b", "c"])) - def test_set_coords_multiindex_level(self): + def test_set_coords_multiindex_level(self) -> None: with pytest.raises( ValueError, match=r"cannot set or update variable.*corrupt.*index " ): self.mda["level_1"] = range(4) - def test_coords_replacement_alignment(self): + def test_coords_replacement_alignment(self) -> None: # regression test for GH725 arr = DataArray([0, 1, 2], dims=["abc"]) new_coord = DataArray([1, 2, 3], dims=["abc"], coords=[[1, 2, 3]]) @@ -1489,25 +1512,25 @@ def test_coords_replacement_alignment(self): expected = DataArray([0, 1, 2], coords=[("abc", [1, 2, 3])]) assert_identical(arr, expected) - def test_coords_non_string(self): + def test_coords_non_string(self) -> None: arr = DataArray(0, coords={1: 2}) actual = arr.coords[1] expected = DataArray(2, coords={1: 2}, name=1) assert_identical(actual, expected) - def test_coords_delitem_delete_indexes(self): + def test_coords_delitem_delete_indexes(self) -> None: # regression test for GH3746 arr = DataArray(np.ones((2,)), dims="x", coords={"x": [0, 1]}) del arr.coords["x"] assert "x" not in arr.xindexes - def test_coords_delitem_multiindex_level(self): + def test_coords_delitem_multiindex_level(self) -> None: with pytest.raises( ValueError, match=r"cannot remove coordinate.*corrupt.*index " ): del self.mda.coords["level_1"] - def test_broadcast_like(self): + def test_broadcast_like(self) -> None: arr1 = DataArray( np.ones((2, 3)), dims=["x", "y"], @@ -1532,7 +1555,7 @@ def test_broadcast_like(self): assert_identical(orig3.broadcast_like(orig4), new3.transpose("y", "x")) assert_identical(orig4.broadcast_like(orig3), new4) - def test_reindex_like(self): + def test_reindex_like(self) -> None: foo = DataArray(np.random.randn(5, 6), [("x", range(5)), ("y", range(6))]) bar = foo[:2, :2] assert_identical(foo.reindex_like(bar), bar) @@ -1542,7 +1565,7 @@ def test_reindex_like(self): expected[:2, :2] = bar assert_identical(bar.reindex_like(foo), expected) - def test_reindex_like_no_index(self): + def test_reindex_like_no_index(self) -> None: foo = DataArray(np.random.randn(5, 6), dims=["x", "y"]) assert_identical(foo, foo.reindex_like(foo)) @@ -1550,15 +1573,15 @@ def test_reindex_like_no_index(self): with pytest.raises(ValueError, match=r"different size for unlabeled"): foo.reindex_like(bar) - def test_reindex_regressions(self): + def test_reindex_regressions(self) -> None: da = DataArray(np.random.randn(5), coords=[("time", range(5))]) time2 = DataArray(np.arange(5), dims="time2") with pytest.raises(ValueError): da.reindex(time=time2) # regression test for #736, reindex can not change complex nums dtype - x = np.array([1, 2, 3], dtype=complex) - x = DataArray(x, coords=[[0.1, 0.2, 0.3]]) + xnp = np.array([1, 2, 3], dtype=complex) + x = DataArray(xnp, coords=[[0.1, 0.2, 0.3]]) y = DataArray([2, 5, 6, 7, 8], coords=[[-1.1, 0.21, 0.31, 0.41, 0.51]]) re_dtype = x.reindex_like(y, method="pad").dtype assert x.dtype == re_dtype @@ -1580,7 +1603,7 @@ def test_reindex_method(self) -> None: assert_identical(expected, actual) @pytest.mark.parametrize("fill_value", [dtypes.NA, 2, 2.0, {None: 2, "u": 1}]) - def test_reindex_fill_value(self, fill_value): + def test_reindex_fill_value(self, fill_value) -> None: x = DataArray([10, 20], dims="y", coords={"y": [0, 1], "u": ("y", [1, 2])}) y = [0, 1, 2] if fill_value == dtypes.NA: @@ -1601,7 +1624,7 @@ def test_reindex_fill_value(self, fill_value): assert_identical(expected, actual) @pytest.mark.parametrize("dtype", [str, bytes]) - def test_reindex_str_dtype(self, dtype): + def test_reindex_str_dtype(self, dtype) -> None: data = DataArray( [1, 2], dims="x", coords={"x": np.array(["a", "b"], dtype=dtype)} @@ -1613,7 +1636,7 @@ def test_reindex_str_dtype(self, dtype): assert_identical(expected, actual) assert actual.dtype == expected.dtype - def test_rename(self): + def test_rename(self) -> None: renamed = self.dv.rename("bar") assert_identical(renamed.to_dataset(), self.ds.rename({"foo": "bar"})) assert renamed.name == "bar" @@ -1626,7 +1649,7 @@ def test_rename(self): renamed_kwargs = self.dv.x.rename(x="z").rename("z") assert_identical(renamed, renamed_kwargs) - def test_init_value(self): + def test_init_value(self) -> None: expected = DataArray( np.full((3, 4), 3), dims=["x", "y"], coords=[range(3), range(4)] ) @@ -1652,7 +1675,7 @@ def test_init_value(self): with pytest.raises(ValueError, match=r"does not match the 0 dim"): DataArray(np.array(1), coords=[("x", np.arange(10))]) - def test_swap_dims(self): + def test_swap_dims(self) -> None: array = DataArray(np.random.randn(3), {"x": list("abc")}, "x") expected = DataArray(array.values, {"x": ("y", list("abc"))}, dims="y") actual = array.swap_dims({"x": "y"}) @@ -1677,7 +1700,7 @@ def test_swap_dims(self): for dim_name in set().union(expected.xindexes.keys(), actual.xindexes.keys()): assert actual.xindexes[dim_name].equals(expected.xindexes[dim_name]) - def test_expand_dims_error(self): + def test_expand_dims_error(self) -> None: array = DataArray( np.random.randn(3, 4), dims=["x", "dim_0"], @@ -1685,7 +1708,7 @@ def test_expand_dims_error(self): attrs={"key": "entry"}, ) - with pytest.raises(TypeError, match=r"dim should be hashable or"): + with pytest.raises(TypeError, match=r"dim should be Hashable or"): array.expand_dims(0) with pytest.raises(ValueError, match=r"lengths of dim and axis"): # dims and axis argument should be the same length @@ -1723,7 +1746,7 @@ def test_expand_dims_error(self): with pytest.raises(ValueError): array.expand_dims({"d": 4}, e=4) - def test_expand_dims(self): + def test_expand_dims(self) -> None: array = DataArray( np.random.randn(3, 4), dims=["x", "dim_0"], @@ -1781,7 +1804,7 @@ def test_expand_dims(self): roundtripped = actual.squeeze(["y", "z"], drop=True) assert_identical(array, roundtripped) - def test_expand_dims_with_scalar_coordinate(self): + def test_expand_dims_with_scalar_coordinate(self) -> None: array = DataArray( np.random.randn(3, 4), dims=["x", "dim_0"], @@ -1799,7 +1822,7 @@ def test_expand_dims_with_scalar_coordinate(self): roundtripped = actual.squeeze(["z"], drop=False) assert_identical(array, roundtripped) - def test_expand_dims_with_greater_dim_size(self): + def test_expand_dims_with_greater_dim_size(self) -> None: array = DataArray( np.random.randn(3, 4), dims=["x", "dim_0"], @@ -1840,7 +1863,7 @@ def test_expand_dims_with_greater_dim_size(self): ).drop_vars("dim_0") assert_identical(other_way_expected, other_way) - def test_set_index(self): + def test_set_index(self) -> None: indexes = [self.mindex.get_level_values(n) for n in self.mindex.names] coords = {idx.name: ("x", idx) for idx in indexes} array = DataArray(self.mda.values, coords=coords, dims="x") @@ -1871,7 +1894,7 @@ def test_set_index(self): with pytest.raises(ValueError, match=r".*variable\(s\) do not exist"): obj.set_index(x="level_4") - def test_reset_index(self): + def test_reset_index(self) -> None: indexes = [self.mindex.get_level_values(n) for n in self.mindex.names] coords = {idx.name: ("x", idx) for idx in indexes} coords["x"] = ("x", self.mindex.values) @@ -1908,14 +1931,14 @@ def test_reset_index(self): assert_identical(obj, array, check_default_indexes=False) assert len(obj.xindexes) == 0 - def test_reset_index_keep_attrs(self): + def test_reset_index_keep_attrs(self) -> None: coord_1 = DataArray([1, 2], dims=["coord_1"], attrs={"attrs": True}) da = DataArray([1, 0], [coord_1]) obj = da.reset_index("coord_1") assert_identical(obj, da, check_default_indexes=False) assert len(obj.xindexes) == 0 - def test_reorder_levels(self): + def test_reorder_levels(self) -> None: midx = self.mindex.reorder_levels(["level_2", "level_1"]) expected = DataArray(self.mda.values, coords={"x": midx}, dims="x") @@ -1930,11 +1953,11 @@ def test_reorder_levels(self): with pytest.raises(ValueError, match=r"has no MultiIndex"): array.reorder_levels(x=["level_1", "level_2"]) - def test_dataset_getitem(self): + def test_dataset_getitem(self) -> None: dv = self.ds["foo"] assert_identical(dv, self.dv) - def test_array_interface(self): + def test_array_interface(self) -> None: assert_array_equal(np.asarray(self.dv), self.x) # test patched in methods assert_array_equal(self.dv.astype(float), self.v.astype(float)) @@ -1948,27 +1971,27 @@ def test_array_interface(self): bar = Variable(["x", "y"], np.zeros((10, 20))) assert_equal(self.dv, np.maximum(self.dv, bar)) - def test_astype_attrs(self): + def test_astype_attrs(self) -> None: for v in [self.va.copy(), self.mda.copy(), self.ds.copy()]: v.attrs["foo"] = "bar" assert v.attrs == v.astype(float).attrs assert not v.astype(float, keep_attrs=False).attrs - def test_astype_dtype(self): + def test_astype_dtype(self) -> None: original = DataArray([-1, 1, 2, 3, 1000]) converted = original.astype(float) assert_array_equal(original, converted) assert np.issubdtype(original.dtype, np.integer) assert np.issubdtype(converted.dtype, np.floating) - def test_astype_order(self): + def test_astype_order(self) -> None: original = DataArray([[1, 2], [3, 4]]) converted = original.astype("d", order="F") assert_equal(original, converted) assert original.values.flags["C_CONTIGUOUS"] assert converted.values.flags["F_CONTIGUOUS"] - def test_astype_subok(self): + def test_astype_subok(self) -> None: class NdArraySubclass(np.ndarray): pass @@ -1981,7 +2004,7 @@ class NdArraySubclass(np.ndarray): assert not isinstance(converted_not_subok.data, NdArraySubclass) assert isinstance(converted_subok.data, NdArraySubclass) - def test_is_null(self): + def test_is_null(self) -> None: x = np.random.RandomState(42).randn(5, 6) x[x < 0] = np.nan original = DataArray(x, [-np.arange(5), np.arange(6)], ["x", "y"]) @@ -1989,7 +2012,7 @@ def test_is_null(self): assert_identical(expected, original.isnull()) assert_identical(~expected, original.notnull()) - def test_math(self): + def test_math(self) -> None: x = self.x v = self.v a = self.dv @@ -2005,25 +2028,25 @@ def test_math(self): assert_equal(a, a + 0 * a) assert_equal(a, 0 * a + a) - def test_math_automatic_alignment(self): + def test_math_automatic_alignment(self) -> None: a = DataArray(range(5), [("x", range(5))]) b = DataArray(range(5), [("x", range(1, 6))]) expected = DataArray(np.ones(4), [("x", [1, 2, 3, 4])]) assert_identical(a - b, expected) - def test_non_overlapping_dataarrays_return_empty_result(self): + def test_non_overlapping_dataarrays_return_empty_result(self) -> None: a = DataArray(range(5), [("x", range(5))]) result = a.isel(x=slice(2)) + a.isel(x=slice(2, None)) assert len(result["x"]) == 0 - def test_empty_dataarrays_return_empty_result(self): + def test_empty_dataarrays_return_empty_result(self) -> None: a = DataArray(data=[]) result = a * a assert len(result["dim_0"]) == 0 - def test_inplace_math_basics(self): + def test_inplace_math_basics(self) -> None: x = self.x a = self.dv v = a.variable @@ -2034,7 +2057,7 @@ def test_inplace_math_basics(self): assert_array_equal(b.values, x) assert source_ndarray(b.values) is x - def test_inplace_math_error(self): + def test_inplace_math_error(self) -> None: data = np.random.rand(4) times = np.arange(4) foo = DataArray(data, coords=[times], dims=["time"]) @@ -2046,7 +2069,7 @@ def test_inplace_math_error(self): # Check error throwing prevented inplace operation assert_array_equal(foo.coords["time"], b) - def test_inplace_math_automatic_alignment(self): + def test_inplace_math_automatic_alignment(self) -> None: a = DataArray(range(5), [("x", range(5))]) b = DataArray(range(1, 6), [("x", range(1, 6))]) with pytest.raises(xr.MergeError, match="Automatic alignment is not supported"): @@ -2054,7 +2077,7 @@ def test_inplace_math_automatic_alignment(self): with pytest.raises(xr.MergeError, match="Automatic alignment is not supported"): b += a - def test_math_name(self): + def test_math_name(self) -> None: # Verify that name is preserved only when it can be done unambiguously. # The rule (copied from pandas.Series) is keep the current name only if # the other object has the same name or no name attribute and this @@ -2069,7 +2092,7 @@ def test_math_name(self): assert (a["x"] + 0).name == "x" assert (a + a["x"]).name is None - def test_math_with_coords(self): + def test_math_with_coords(self) -> None: coords = { "x": [-1, -2], "y": ["ab", "cd", "ef"], @@ -2124,7 +2147,7 @@ def test_math_with_coords(self): actual = alt + orig assert_identical(expected, actual) - def test_index_math(self): + def test_index_math(self) -> None: orig = DataArray(range(3), dims="x", name="x") actual = orig + 1 expected = DataArray(1 + np.arange(3), dims="x", name="x") @@ -2138,20 +2161,20 @@ def test_index_math(self): actual = orig > orig[0] assert_identical(expected, actual) - def test_dataset_math(self): + def test_dataset_math(self) -> None: # more comprehensive tests with multiple dataset variables obs = Dataset( {"tmin": ("x", np.arange(5)), "tmax": ("x", 10 + np.arange(5))}, {"x": ("x", 0.5 * np.arange(5)), "loc": ("x", range(-2, 3))}, ) - actual = 2 * obs["tmax"] - expected = DataArray(2 * (10 + np.arange(5)), obs.coords, name="tmax") - assert_identical(actual, expected) + actual1 = 2 * obs["tmax"] + expected1 = DataArray(2 * (10 + np.arange(5)), obs.coords, name="tmax") + assert_identical(actual1, expected1) - actual = obs["tmax"] - obs["tmin"] - expected = DataArray(10 * np.ones(5), obs.coords) - assert_identical(actual, expected) + actual2 = obs["tmax"] - obs["tmin"] + expected2 = DataArray(10 * np.ones(5), obs.coords) + assert_identical(actual2, expected2) sim = Dataset( { @@ -2162,29 +2185,29 @@ def test_dataset_math(self): } ) - actual = sim["tmin"] - obs["tmin"] - expected = DataArray(np.ones(5), obs.coords, name="tmin") - assert_identical(actual, expected) + actual3 = sim["tmin"] - obs["tmin"] + expected3 = DataArray(np.ones(5), obs.coords, name="tmin") + assert_identical(actual3, expected3) - actual = -obs["tmin"] + sim["tmin"] - assert_identical(actual, expected) + actual4 = -obs["tmin"] + sim["tmin"] + assert_identical(actual4, expected3) - actual = sim["tmin"].copy() - actual -= obs["tmin"] - assert_identical(actual, expected) + actual5 = sim["tmin"].copy() + actual5 -= obs["tmin"] + assert_identical(actual5, expected3) - actual = sim.copy() - actual["tmin"] = sim["tmin"] - obs["tmin"] - expected = Dataset( + actual6 = sim.copy() + actual6["tmin"] = sim["tmin"] - obs["tmin"] + expected6 = Dataset( {"tmin": ("x", np.ones(5)), "tmax": ("x", sim["tmax"].values)}, obs.coords ) - assert_identical(actual, expected) + assert_identical(actual6, expected6) - actual = sim.copy() - actual["tmin"] -= obs["tmin"] - assert_identical(actual, expected) + actual7 = sim.copy() + actual7["tmin"] -= obs["tmin"] + assert_identical(actual7, expected6) - def test_stack_unstack(self): + def test_stack_unstack(self) -> None: orig = DataArray( [[0, 1], [2, 3]], dims=["x", "y"], @@ -2224,7 +2247,7 @@ def test_stack_unstack(self): unstacked = stacked.unstack() assert_identical(orig, unstacked.transpose(*dims)) - def test_stack_unstack_decreasing_coordinate(self): + def test_stack_unstack_decreasing_coordinate(self) -> None: # regression test for GH980 orig = DataArray( np.random.rand(3, 4), @@ -2235,25 +2258,25 @@ def test_stack_unstack_decreasing_coordinate(self): actual = stacked.unstack("allpoints") assert_identical(orig, actual) - def test_unstack_pandas_consistency(self): + def test_unstack_pandas_consistency(self) -> None: df = pd.DataFrame({"foo": range(3), "x": ["a", "b", "b"], "y": [0, 0, 1]}) s = df.set_index(["x", "y"])["foo"] expected = DataArray(s.unstack(), name="foo") actual = DataArray(s, dims="z").unstack("z") assert_identical(expected, actual) - def test_stack_nonunique_consistency(self, da): + def test_stack_nonunique_consistency(self, da) -> None: da = da.isel(time=0, drop=True) # 2D actual = da.stack(z=["a", "x"]) expected = DataArray(da.to_pandas().stack(), dims="z") assert_identical(expected, actual) - def test_to_unstacked_dataset_raises_value_error(self): + def test_to_unstacked_dataset_raises_value_error(self) -> None: data = DataArray([0, 1], dims="x", coords={"x": [0, 1]}) with pytest.raises(ValueError, match="'x' is not a stacked coordinate"): data.to_unstacked_dataset("x", 0) - def test_transpose(self): + def test_transpose(self) -> None: da = DataArray( np.random.randn(3, 4, 5), dims=("x", "y", "z"), @@ -2301,10 +2324,10 @@ def test_transpose(self): with pytest.warns(UserWarning): da.transpose("not_a_dim", "y", "x", ..., missing_dims="warn") - def test_squeeze(self): + def test_squeeze(self) -> None: assert_equal(self.dv.variable.squeeze(), self.dv.squeeze().variable) - def test_squeeze_drop(self): + def test_squeeze_drop(self) -> None: array = DataArray([1], [("x", [0])]) expected = DataArray(1) actual = array.squeeze(drop=True) @@ -2328,7 +2351,7 @@ def test_squeeze_drop(self): with pytest.raises(ValueError): array.squeeze(axis=0, dim="dim_1") - def test_drop_coordinates(self): + def test_drop_coordinates(self) -> None: expected = DataArray(np.random.randn(2, 3), dims=["x", "y"]) arr = expected.copy() arr.coords["z"] = 2 @@ -2354,20 +2377,21 @@ def test_drop_coordinates(self): actual = renamed.drop_vars("foo", errors="ignore") assert_identical(actual, renamed) - def test_drop_multiindex_level(self): - with pytest.raises( - ValueError, match=r"cannot remove coordinate.*corrupt.*index " - ): - self.mda.drop_vars("level_1") + def test_drop_multiindex_level(self) -> None: + # GH6505 + expected = self.mda.drop_vars(["x", "level_1", "level_2"]) + with pytest.warns(DeprecationWarning): + actual = self.mda.drop_vars("level_1") + assert_identical(expected, actual) - def test_drop_all_multiindex_levels(self): + def test_drop_all_multiindex_levels(self) -> None: dim_levels = ["x", "level_1", "level_2"] actual = self.mda.drop_vars(dim_levels) # no error, multi-index dropped for key in dim_levels: assert key not in actual.xindexes - def test_drop_index_labels(self): + def test_drop_index_labels(self) -> None: arr = DataArray(np.random.randn(2, 3), coords={"y": [0, 1, 2]}, dims=["x", "y"]) actual = arr.drop_sel(y=[0, 1]) expected = arr[:, 2:] @@ -2380,15 +2404,15 @@ def test_drop_index_labels(self): assert_identical(actual, expected) with pytest.warns(DeprecationWarning): - arr.drop([0, 1, 3], dim="y", errors="ignore") + arr.drop([0, 1, 3], dim="y", errors="ignore") # type: ignore - def test_drop_index_positions(self): + def test_drop_index_positions(self) -> None: arr = DataArray(np.random.randn(2, 3), dims=["x", "y"]) actual = arr.drop_isel(y=[0, 1]) expected = arr[:, 2:] assert_identical(actual, expected) - def test_dropna(self): + def test_dropna(self) -> None: x = np.random.randn(4, 4) x[::2, 0] = np.nan arr = DataArray(x, dims=["a", "b"]) @@ -2407,25 +2431,25 @@ def test_dropna(self): expected = arr[:, 1:] assert_identical(actual, expected) - def test_where(self): + def test_where(self) -> None: arr = DataArray(np.arange(4), dims="x") expected = arr.sel(x=slice(2)) actual = arr.where(arr.x < 2, drop=True) assert_identical(actual, expected) - def test_where_lambda(self): + def test_where_lambda(self) -> None: arr = DataArray(np.arange(4), dims="y") expected = arr.sel(y=slice(2)) actual = arr.where(lambda x: x.y < 2, drop=True) assert_identical(actual, expected) - def test_where_string(self): + def test_where_string(self) -> None: array = DataArray(["a", "b"]) expected = DataArray(np.array(["a", np.nan], dtype=object)) actual = array.where([True, False]) assert_identical(actual, expected) - def test_cumops(self): + def test_cumops(self) -> None: coords = { "x": [-1, -2], "y": ["ab", "cd", "ef"], @@ -2526,7 +2550,7 @@ def test_reduce_keepdims(self) -> None: assert_equal(actual, expected) @requires_bottleneck - def test_reduce_keepdims_bottleneck(self): + def test_reduce_keepdims_bottleneck(self) -> None: import bottleneck coords = { @@ -2542,7 +2566,7 @@ def test_reduce_keepdims_bottleneck(self): expected = orig.mean(keepdims=True) assert_equal(actual, expected) - def test_reduce_dtype(self): + def test_reduce_dtype(self) -> None: coords = { "x": [-1, -2], "y": ["ab", "cd", "ef"], @@ -2554,7 +2578,7 @@ def test_reduce_dtype(self): for dtype in [np.float16, np.float32, np.float64]: assert orig.astype(float).mean(dtype=dtype).dtype == dtype - def test_reduce_out(self): + def test_reduce_out(self) -> None: coords = { "x": [-1, -2], "y": ["ab", "cd", "ef"], @@ -2619,7 +2643,7 @@ def test_quantile_interpolation_deprecated(self, method) -> None: with pytest.raises(TypeError, match="interpolation and method keywords"): da.quantile(q, method=method, interpolation=method) - def test_reduce_keep_attrs(self): + def test_reduce_keep_attrs(self) -> None: # Test dropped attrs vm = self.va.mean() assert len(vm.attrs) == 0 @@ -2630,7 +2654,7 @@ def test_reduce_keep_attrs(self): assert len(vm.attrs) == len(self.attrs) assert vm.attrs == self.attrs - def test_assign_attrs(self): + def test_assign_attrs(self) -> None: expected = DataArray([], attrs=dict(a=1, b=2)) expected.attrs["a"] = 1 expected.attrs["b"] = 2 @@ -2647,7 +2671,7 @@ def test_assign_attrs(self): @pytest.mark.parametrize( "func", [lambda x: x.clip(0, 1), lambda x: np.float64(1.0) * x, np.abs, abs] ) - def test_propagate_attrs(self, func): + def test_propagate_attrs(self, func) -> None: da = DataArray(self.va) # test defaults @@ -2659,7 +2683,7 @@ def test_propagate_attrs(self, func): with set_options(keep_attrs=True): assert func(da).attrs == da.attrs - def test_fillna(self): + def test_fillna(self) -> None: a = DataArray([np.nan, 1, np.nan, 3], coords={"x": range(4)}, dims="x") actual = a.fillna(-1) expected = DataArray([-1, 1, -1, 3], coords={"x": range(4)}, dims="x") @@ -2685,7 +2709,7 @@ def test_fillna(self): with pytest.raises(ValueError, match=r"broadcast"): a.fillna([1, 2]) - def test_align(self): + def test_align(self) -> None: array = DataArray( np.random.random((6, 8)), coords={"x": list("abcdef")}, dims=["x", "y"] ) @@ -2693,7 +2717,7 @@ def test_align(self): assert_identical(array1, array[:5]) assert_identical(array2, array[:5]) - def test_align_dtype(self): + def test_align_dtype(self) -> None: # regression test for #264 x1 = np.arange(30) x2 = np.arange(5, 35) @@ -2702,7 +2726,7 @@ def test_align_dtype(self): c, d = align(a, b, join="outer") assert c.dtype == np.float32 - def test_align_copy(self): + def test_align_copy(self) -> None: x = DataArray([1, 2, 3], coords=[("a", [1, 2, 3])]) y = DataArray([1, 2], coords=[("a", [3, 1])]) @@ -2729,7 +2753,7 @@ def test_align_copy(self): assert_identical(x, x2) assert source_ndarray(x2.data) is not source_ndarray(x.data) - def test_align_override(self): + def test_align_override(self) -> None: left = DataArray([1, 2, 3], dims="x", coords={"x": [0, 1, 2]}) right = DataArray( np.arange(9).reshape((3, 3)), @@ -2777,13 +2801,13 @@ def test_align_override(self): ], ], ) - def test_align_override_error(self, darrays): + def test_align_override_error(self, darrays) -> None: with pytest.raises( ValueError, match=r"cannot align.*join.*override.*same size" ): xr.align(*darrays, join="override") - def test_align_exclude(self): + def test_align_exclude(self) -> None: x = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, -2]), ("b", [3, 4])]) y = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, 20]), ("b", [5, 6])]) z = DataArray([1], dims=["a"], coords={"a": [20], "b": 7}) @@ -2804,7 +2828,7 @@ def test_align_exclude(self): assert_identical(expected_y2, y2) assert_identical(expected_z2, z2) - def test_align_indexes(self): + def test_align_indexes(self) -> None: x = DataArray([1, 2, 3], coords=[("a", [-1, 10, -2])]) y = DataArray([1, 2], coords=[("a", [-2, -1])]) @@ -2818,13 +2842,13 @@ def test_align_indexes(self): expected_x2 = DataArray([3, np.nan, 2, 1], coords=[("a", [-2, 7, 10, -1])]) assert_identical(expected_x2, x2) - def test_align_without_indexes_exclude(self): + def test_align_without_indexes_exclude(self) -> None: arrays = [DataArray([1, 2, 3], dims=["x"]), DataArray([1, 2], dims=["x"])] result0, result1 = align(*arrays, exclude=["x"]) assert_identical(result0, arrays[0]) assert_identical(result1, arrays[1]) - def test_align_mixed_indexes(self): + def test_align_mixed_indexes(self) -> None: array_no_coord = DataArray([1, 2], dims=["x"]) array_with_coord = DataArray([1, 2], coords=[("x", ["a", "b"])]) result0, result1 = align(array_no_coord, array_with_coord) @@ -2835,7 +2859,7 @@ def test_align_mixed_indexes(self): assert_identical(result0, array_no_coord) assert_identical(result1, array_with_coord) - def test_align_without_indexes_errors(self): + def test_align_without_indexes_errors(self) -> None: with pytest.raises( ValueError, match=r"cannot.*align.*dimension.*conflicting.*sizes.*", @@ -2851,7 +2875,7 @@ def test_align_without_indexes_errors(self): DataArray([1, 2], coords=[("x", [0, 1])]), ) - def test_align_str_dtype(self): + def test_align_str_dtype(self) -> None: a = DataArray([0, 1], dims=["x"], coords={"x": ["a", "b"]}) b = DataArray([1, 2], dims=["x"], coords={"x": ["b", "c"]}) @@ -2871,7 +2895,7 @@ def test_align_str_dtype(self): assert_identical(expected_b, actual_b) assert expected_b.x.dtype == actual_b.x.dtype - def test_broadcast_arrays(self): + def test_broadcast_arrays(self) -> None: x = DataArray([1, 2], coords=[("a", [-1, -2])], name="x") y = DataArray([1, 2], coords=[("b", [3, 4])], name="y") x2, y2 = broadcast(x, y) @@ -2889,7 +2913,7 @@ def test_broadcast_arrays(self): assert_identical(expected_x2, x2) assert_identical(expected_y2, y2) - def test_broadcast_arrays_misaligned(self): + def test_broadcast_arrays_misaligned(self) -> None: # broadcast on misaligned coords must auto-align x = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, -2]), ("b", [3, 4])]) y = DataArray([1, 2], coords=[("a", [-1, 20])]) @@ -2905,7 +2929,7 @@ def test_broadcast_arrays_misaligned(self): assert_identical(expected_x2, x2) assert_identical(expected_y2, y2) - def test_broadcast_arrays_nocopy(self): + def test_broadcast_arrays_nocopy(self) -> None: # Test that input data is not copied over in case # no alteration is needed x = DataArray([1, 2], coords=[("a", [-1, -2])], name="x") @@ -2923,7 +2947,7 @@ def test_broadcast_arrays_nocopy(self): assert_identical(x, x2) assert source_ndarray(x2.data) is source_ndarray(x.data) - def test_broadcast_arrays_exclude(self): + def test_broadcast_arrays_exclude(self) -> None: x = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, -2]), ("b", [3, 4])]) y = DataArray([1, 2], coords=[("a", [-1, 20])]) z = DataArray(5, coords={"b": 5}) @@ -2941,7 +2965,7 @@ def test_broadcast_arrays_exclude(self): assert_identical(expected_y2, y2) assert_identical(expected_z2, z2) - def test_broadcast_coordinates(self): + def test_broadcast_coordinates(self) -> None: # regression test for GH649 ds = Dataset({"a": (["x", "y"], np.ones((5, 6)))}) x_bc, y_bc, a_bc = broadcast(ds.x, ds.y, ds.a) @@ -2953,7 +2977,7 @@ def test_broadcast_coordinates(self): assert_identical(exp_x, x_bc) assert_identical(exp_y, y_bc) - def test_to_pandas(self): + def test_to_pandas(self) -> None: # 0d actual = DataArray(42).to_pandas() expected = np.array(42) @@ -2985,10 +3009,10 @@ def test_to_pandas(self): roundtripped = DataArray(da.to_pandas()).drop_vars(dims) assert_identical(da, roundtripped) - with pytest.raises(ValueError, match=r"cannot convert"): + with pytest.raises(ValueError, match=r"Cannot convert"): DataArray(np.random.randn(1, 2, 3, 4, 5)).to_pandas() - def test_to_dataframe(self): + def test_to_dataframe(self) -> None: # regression test for #260 arr_np = np.random.randn(3, 4) @@ -3022,7 +3046,7 @@ def test_to_dataframe(self): with pytest.raises(ValueError, match=r"unnamed"): arr.to_dataframe() - def test_to_dataframe_multiindex(self): + def test_to_dataframe_multiindex(self) -> None: # regression test for #3008 arr_np = np.random.randn(4, 3) @@ -3037,7 +3061,7 @@ def test_to_dataframe_multiindex(self): assert_array_equal(actual.index.levels[1], ["a", "b"]) assert_array_equal(actual.index.levels[2], [5, 6, 7]) - def test_to_dataframe_0length(self): + def test_to_dataframe_0length(self) -> None: # regression test for #3008 arr_np = np.random.randn(4, 0) @@ -3049,7 +3073,7 @@ def test_to_dataframe_0length(self): assert len(actual) == 0 assert_array_equal(actual.index.names, list("ABC")) - def test_to_pandas_name_matches_coordinate(self): + def test_to_pandas_name_matches_coordinate(self) -> None: # coordinate with same name as array arr = DataArray([1, 2, 3], dims="x", name="x") series = arr.to_series() @@ -3062,7 +3086,7 @@ def test_to_pandas_name_matches_coordinate(self): expected = series.to_frame() assert expected.equals(frame) - def test_to_and_from_series(self): + def test_to_and_from_series(self) -> None: expected = self.dv.to_dataframe()["foo"] actual = self.dv.to_series() assert_array_equal(expected.values, actual.values) @@ -3077,7 +3101,7 @@ def test_to_and_from_series(self): expected_da, DataArray.from_series(actual).drop_vars(["x", "y"]) ) - def test_from_series_multiindex(self): + def test_from_series_multiindex(self) -> None: # GH:3951 df = pd.DataFrame({"B": [1, 2, 3], "A": [4, 5, 6]}) df = df.rename_axis("num").rename_axis("alpha", axis=1) @@ -3086,7 +3110,7 @@ def test_from_series_multiindex(self): assert (actual.sel(alpha="A") == [4, 5, 6]).all() @requires_sparse - def test_from_series_sparse(self): + def test_from_series_sparse(self) -> None: import sparse series = pd.Series([1, 2], index=[("a", 1), ("b", 2)]) @@ -3099,7 +3123,7 @@ def test_from_series_sparse(self): assert_identical(actual_sparse, actual_dense) @requires_sparse - def test_from_multiindex_series_sparse(self): + def test_from_multiindex_series_sparse(self) -> None: # regression test for GH4019 import sparse @@ -3116,7 +3140,7 @@ def test_from_multiindex_series_sparse(self): np.testing.assert_equal(actual_coords, expected_coords) - def test_to_and_from_empty_series(self): + def test_to_and_from_empty_series(self) -> None: # GH697 expected = pd.Series([], dtype=np.float64) da = DataArray.from_series(expected) @@ -3125,7 +3149,7 @@ def test_to_and_from_empty_series(self): assert len(actual) == 0 assert expected.equals(actual) - def test_series_categorical_index(self): + def test_series_categorical_index(self) -> None: # regression test for GH700 if not hasattr(pd, "CategoricalIndex"): pytest.skip("requires pandas with CategoricalIndex") @@ -3134,10 +3158,12 @@ def test_series_categorical_index(self): arr = DataArray(s) assert "'a'" in repr(arr) # should not error - def test_to_and_from_dict(self): + @pytest.mark.parametrize("encoding", [True, False]) + def test_to_and_from_dict(self, encoding) -> None: array = DataArray( np.random.randn(2, 3), {"x": ["a", "b"]}, ["x", "y"], name="foo" ) + array.encoding = {"bar": "spam"} expected = { "name": "foo", "dims": ("x", "y"), @@ -3145,7 +3171,9 @@ def test_to_and_from_dict(self): "attrs": {}, "coords": {"x": {"dims": ("x",), "data": ["a", "b"], "attrs": {}}}, } - actual = array.to_dict() + if encoding: + expected["encoding"] = {"bar": "spam"} + actual = array.to_dict(encoding=encoding) # check that they are identical assert expected == actual @@ -3192,10 +3220,10 @@ def test_to_and_from_dict(self): endiantype = "U1" expected_no_data["coords"]["x"].update({"dtype": endiantype, "shape": (2,)}) expected_no_data.update({"dtype": "float64", "shape": (2, 3)}) - actual_no_data = array.to_dict(data=False) + actual_no_data = array.to_dict(data=False, encoding=encoding) assert expected_no_data == actual_no_data - def test_to_and_from_dict_with_time_dim(self): + def test_to_and_from_dict_with_time_dim(self) -> None: x = np.random.randn(10, 3) t = pd.date_range("20130101", periods=10) lat = [77.7, 83.2, 76] @@ -3203,7 +3231,7 @@ def test_to_and_from_dict_with_time_dim(self): roundtripped = DataArray.from_dict(da.to_dict()) assert_identical(da, roundtripped) - def test_to_and_from_dict_with_nan_nat(self): + def test_to_and_from_dict_with_nan_nat(self) -> None: y = np.random.randn(10, 3) y[2] = np.nan t = pd.Series(pd.date_range("20130101", periods=10)) @@ -3213,7 +3241,7 @@ def test_to_and_from_dict_with_nan_nat(self): roundtripped = DataArray.from_dict(da.to_dict()) assert_identical(da, roundtripped) - def test_to_dict_with_numpy_attrs(self): + def test_to_dict_with_numpy_attrs(self) -> None: # this doesn't need to roundtrip x = np.random.randn(10, 3) t = list("abcdefghij") @@ -3225,8 +3253,8 @@ def test_to_dict_with_numpy_attrs(self): } da = DataArray(x, {"t": t, "lat": lat}, dims=["t", "lat"], attrs=attrs) expected_attrs = { - "created": attrs["created"].item(), - "coords": attrs["coords"].tolist(), + "created": attrs["created"].item(), # type: ignore[attr-defined] + "coords": attrs["coords"].tolist(), # type: ignore[attr-defined] "maintainer": "bar", } actual = da.to_dict() @@ -3234,7 +3262,7 @@ def test_to_dict_with_numpy_attrs(self): # check that they are identical assert expected_attrs == actual["attrs"] - def test_to_masked_array(self): + def test_to_masked_array(self) -> None: rs = np.random.RandomState(44) x = rs.random_sample(size=(10, 20)) x_masked = np.ma.masked_where(x < 0.5, x) @@ -3276,7 +3304,7 @@ def test_to_masked_array(self): ma = da.to_masked_array() assert len(ma.mask) == N - def test_to_and_from_cdms2_classic(self): + def test_to_and_from_cdms2_classic(self) -> None: """Classic with 1D axes""" pytest.importorskip("cdms2") @@ -3315,7 +3343,7 @@ def test_to_and_from_cdms2_classic(self): for coord_name in original.coords.keys(): assert_array_equal(original.coords[coord_name], back.coords[coord_name]) - def test_to_and_from_cdms2_sgrid(self): + def test_to_and_from_cdms2_sgrid(self) -> None: """Curvilinear (structured) grid The rectangular grid case is covered by the classic case @@ -3344,7 +3372,7 @@ def test_to_and_from_cdms2_sgrid(self): assert_array_equal(original.coords["lat"], back.coords["lat"]) assert_array_equal(original.coords["lon"], back.coords["lon"]) - def test_to_and_from_cdms2_ugrid(self): + def test_to_and_from_cdms2_ugrid(self) -> None: """Unstructured grid""" pytest.importorskip("cdms2") @@ -3365,7 +3393,7 @@ def test_to_and_from_cdms2_ugrid(self): assert_array_equal(original.coords["lat"], back.coords["lat"]) assert_array_equal(original.coords["lon"], back.coords["lon"]) - def test_to_dataset_whole(self): + def test_to_dataset_whole(self) -> None: unnamed = DataArray([1, 2], dims="x") with pytest.raises(ValueError, match=r"unable to convert unnamed"): unnamed.to_dataset() @@ -3389,7 +3417,7 @@ def test_to_dataset_whole(self): with pytest.raises(TypeError): actual = named.to_dataset("bar") - def test_to_dataset_split(self): + def test_to_dataset_split(self) -> None: array = DataArray([1, 2, 3], coords=[("x", list("abc"))], attrs={"a": 1}) expected = Dataset({"a": 1, "b": 2, "c": 3}, attrs={"a": 1}) actual = array.to_dataset("x") @@ -3406,7 +3434,7 @@ def test_to_dataset_split(self): actual = array.to_dataset("x") assert_identical(expected, actual) - def test_to_dataset_retains_keys(self): + def test_to_dataset_retains_keys(self) -> None: # use dates as convenient non-str objects. Not a specific date test import datetime @@ -3420,7 +3448,7 @@ def test_to_dataset_retains_keys(self): assert_equal(array, result) - def test__title_for_slice(self): + def test__title_for_slice(self) -> None: array = DataArray( np.ones((4, 3, 2)), dims=["a", "b", "c"], @@ -3434,7 +3462,7 @@ def test__title_for_slice(self): a2 = DataArray(np.ones((4, 1)), dims=["a", "b"]) assert "" == a2._title_for_slice() - def test__title_for_slice_truncate(self): + def test__title_for_slice_truncate(self) -> None: array = DataArray(np.ones(4)) array.coords["a"] = "a" * 100 array.coords["b"] = "b" * 100 @@ -3445,13 +3473,13 @@ def test__title_for_slice_truncate(self): assert nchar == len(title) assert title.endswith("...") - def test_dataarray_diff_n1(self): + def test_dataarray_diff_n1(self) -> None: da = DataArray(np.random.randn(3, 4), dims=["x", "y"]) actual = da.diff("y") expected = DataArray(np.diff(da.values, axis=1), dims=["x", "y"]) assert_equal(expected, actual) - def test_coordinate_diff(self): + def test_coordinate_diff(self) -> None: # regression test for GH634 arr = DataArray(range(0, 20, 2), dims=["lon"], coords=[range(10)]) lon = arr.coords["lon"] @@ -3461,7 +3489,7 @@ def test_coordinate_diff(self): @pytest.mark.parametrize("offset", [-5, 0, 1, 2]) @pytest.mark.parametrize("fill_value, dtype", [(2, int), (dtypes.NA, float)]) - def test_shift(self, offset, fill_value, dtype): + def test_shift(self, offset, fill_value, dtype) -> None: arr = DataArray([1, 2, 3], dims="x") actual = arr.shift(x=1, fill_value=fill_value) if fill_value == dtypes.NA: @@ -3477,19 +3505,19 @@ def test_shift(self, offset, fill_value, dtype): actual = arr.shift(x=offset) assert_identical(expected, actual) - def test_roll_coords(self): + def test_roll_coords(self) -> None: arr = DataArray([1, 2, 3], coords={"x": range(3)}, dims="x") actual = arr.roll(x=1, roll_coords=True) expected = DataArray([3, 1, 2], coords=[("x", [2, 0, 1])]) assert_identical(expected, actual) - def test_roll_no_coords(self): + def test_roll_no_coords(self) -> None: arr = DataArray([1, 2, 3], coords={"x": range(3)}, dims="x") actual = arr.roll(x=1) expected = DataArray([3, 1, 2], coords=[("x", [0, 1, 2])]) assert_identical(expected, actual) - def test_copy_with_data(self): + def test_copy_with_data(self) -> None: orig = DataArray( np.random.random(size=(2, 2)), dims=("x", "y"), @@ -3525,7 +3553,7 @@ def test_copy_with_data(self): ], ], ) - def test_copy_coords(self, deep, expected_orig): + def test_copy_coords(self, deep, expected_orig) -> None: """The test fails for the shallow copy, and apparently only on Windows for some reason. In windows coords seem to be immutable unless it's one dataarray deep copied from another.""" @@ -3546,12 +3574,12 @@ def test_copy_coords(self, deep, expected_orig): assert_identical(da["a"], expected_orig) - def test_real_and_imag(self): + def test_real_and_imag(self) -> None: array = DataArray(1 + 2j) assert_identical(array.real, DataArray(1)) assert_identical(array.imag, DataArray(2)) - def test_setattr_raises(self): + def test_setattr_raises(self) -> None: array = DataArray(0, coords={"scalar": 1}, attrs={"foo": "bar"}) with pytest.raises(AttributeError, match=r"cannot set attr"): array.scalar = 2 @@ -3560,7 +3588,7 @@ def test_setattr_raises(self): with pytest.raises(AttributeError, match=r"cannot set attr"): array.other = 2 - def test_full_like(self): + def test_full_like(self) -> None: # For more thorough tests, see test_variable.py da = DataArray( np.random.random(size=(2, 2)), @@ -3572,65 +3600,65 @@ def test_full_like(self): actual = full_like(da, 2) expect = da.copy(deep=True) - expect.values = [[2.0, 2.0], [2.0, 2.0]] + expect.values = np.array([[2.0, 2.0], [2.0, 2.0]]) assert_identical(expect, actual) # override dtype actual = full_like(da, fill_value=True, dtype=bool) - expect.values = [[True, True], [True, True]] + expect.values = np.array([[True, True], [True, True]]) assert expect.dtype == bool assert_identical(expect, actual) with pytest.raises(ValueError, match="'dtype' cannot be dict-like"): full_like(da, fill_value=True, dtype={"x": bool}) - def test_dot(self): + def test_dot(self) -> None: x = np.linspace(-3, 3, 6) y = np.linspace(-3, 3, 5) z = range(4) da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4)) da = DataArray(da_vals, coords=[x, y, z], dims=["x", "y", "z"]) - dm_vals = range(4) - dm = DataArray(dm_vals, coords=[z], dims=["z"]) + dm_vals1 = range(4) + dm1 = DataArray(dm_vals1, coords=[z], dims=["z"]) # nd dot 1d - actual = da.dot(dm) - expected_vals = np.tensordot(da_vals, dm_vals, [2, 0]) - expected = DataArray(expected_vals, coords=[x, y], dims=["x", "y"]) - assert_equal(expected, actual) + actual1 = da.dot(dm1) + expected_vals1 = np.tensordot(da_vals, dm_vals1, (2, 0)) + expected1 = DataArray(expected_vals1, coords=[x, y], dims=["x", "y"]) + assert_equal(expected1, actual1) # all shared dims - actual = da.dot(da) - expected_vals = np.tensordot(da_vals, da_vals, axes=([0, 1, 2], [0, 1, 2])) - expected = DataArray(expected_vals) - assert_equal(expected, actual) + actual2 = da.dot(da) + expected_vals2 = np.tensordot(da_vals, da_vals, axes=([0, 1, 2], [0, 1, 2])) + expected2 = DataArray(expected_vals2) + assert_equal(expected2, actual2) # multiple shared dims - dm_vals = np.arange(20 * 5 * 4).reshape((20, 5, 4)) + dm_vals3 = np.arange(20 * 5 * 4).reshape((20, 5, 4)) j = np.linspace(-3, 3, 20) - dm = DataArray(dm_vals, coords=[j, y, z], dims=["j", "y", "z"]) - actual = da.dot(dm) - expected_vals = np.tensordot(da_vals, dm_vals, axes=([1, 2], [1, 2])) - expected = DataArray(expected_vals, coords=[x, j], dims=["x", "j"]) - assert_equal(expected, actual) + dm3 = DataArray(dm_vals3, coords=[j, y, z], dims=["j", "y", "z"]) + actual3 = da.dot(dm3) + expected_vals3 = np.tensordot(da_vals, dm_vals3, axes=([1, 2], [1, 2])) + expected3 = DataArray(expected_vals3, coords=[x, j], dims=["x", "j"]) + assert_equal(expected3, actual3) # Ellipsis: all dims are shared - actual = da.dot(da, dims=...) - expected = da.dot(da) - assert_equal(expected, actual) + actual4 = da.dot(da, dims=...) + expected4 = da.dot(da) + assert_equal(expected4, actual4) # Ellipsis: not all dims are shared - actual = da.dot(dm, dims=...) - expected = da.dot(dm, dims=("j", "x", "y", "z")) - assert_equal(expected, actual) + actual5 = da.dot(dm3, dims=...) + expected5 = da.dot(dm3, dims=("j", "x", "y", "z")) + assert_equal(expected5, actual5) with pytest.raises(NotImplementedError): - da.dot(dm.to_dataset(name="dm")) + da.dot(dm3.to_dataset(name="dm")) # type: ignore with pytest.raises(TypeError): - da.dot(dm.values) + da.dot(dm3.values) # type: ignore - def test_dot_align_coords(self): + def test_dot_align_coords(self) -> None: # GH 3694 x = np.linspace(-3, 3, 6) @@ -3640,36 +3668,36 @@ def test_dot_align_coords(self): da = DataArray(da_vals, coords=[x, y, z_a], dims=["x", "y", "z"]) z_m = range(2, 6) - dm_vals = range(4) - dm = DataArray(dm_vals, coords=[z_m], dims=["z"]) + dm_vals1 = range(4) + dm1 = DataArray(dm_vals1, coords=[z_m], dims=["z"]) with xr.set_options(arithmetic_join="exact"): with pytest.raises( ValueError, match=r"cannot align.*join.*exact.*not equal.*" ): - da.dot(dm) + da.dot(dm1) - da_aligned, dm_aligned = xr.align(da, dm, join="inner") + da_aligned, dm_aligned = xr.align(da, dm1, join="inner") # nd dot 1d - actual = da.dot(dm) - expected_vals = np.tensordot(da_aligned.values, dm_aligned.values, [2, 0]) - expected = DataArray(expected_vals, coords=[x, da_aligned.y], dims=["x", "y"]) - assert_equal(expected, actual) + actual1 = da.dot(dm1) + expected_vals1 = np.tensordot(da_aligned.values, dm_aligned.values, (2, 0)) + expected1 = DataArray(expected_vals1, coords=[x, da_aligned.y], dims=["x", "y"]) + assert_equal(expected1, actual1) # multiple shared dims - dm_vals = np.arange(20 * 5 * 4).reshape((20, 5, 4)) + dm_vals2 = np.arange(20 * 5 * 4).reshape((20, 5, 4)) j = np.linspace(-3, 3, 20) - dm = DataArray(dm_vals, coords=[j, y, z_m], dims=["j", "y", "z"]) - da_aligned, dm_aligned = xr.align(da, dm, join="inner") - actual = da.dot(dm) - expected_vals = np.tensordot( + dm2 = DataArray(dm_vals2, coords=[j, y, z_m], dims=["j", "y", "z"]) + da_aligned, dm_aligned = xr.align(da, dm2, join="inner") + actual2 = da.dot(dm2) + expected_vals2 = np.tensordot( da_aligned.values, dm_aligned.values, axes=([1, 2], [1, 2]) ) - expected = DataArray(expected_vals, coords=[x, j], dims=["x", "j"]) - assert_equal(expected, actual) + expected2 = DataArray(expected_vals2, coords=[x, j], dims=["x", "j"]) + assert_equal(expected2, actual2) - def test_matmul(self): + def test_matmul(self) -> None: # copied from above (could make a fixture) x = np.linspace(-3, 3, 6) @@ -3682,7 +3710,7 @@ def test_matmul(self): expected = da.dot(da) assert_identical(result, expected) - def test_matmul_align_coords(self): + def test_matmul_align_coords(self) -> None: # GH 3694 x_a = np.arange(6) @@ -3702,7 +3730,7 @@ def test_matmul_align_coords(self): ): da_a @ da_b - def test_binary_op_propagate_indexes(self): + def test_binary_op_propagate_indexes(self) -> None: # regression test for GH2227 self.dv["x"] = np.arange(self.dv.sizes["x"]) expected = self.dv.xindexes["x"] @@ -3713,9 +3741,9 @@ def test_binary_op_propagate_indexes(self): actual = (self.dv > 10).xindexes["x"] assert expected is actual - def test_binary_op_join_setting(self): + def test_binary_op_join_setting(self) -> None: dim = "x" - align_type = "outer" + align_type: Final = "outer" coords_l, coords_r = [0, 1, 2], [1, 2, 3] missing_3 = xr.DataArray(coords_l, [(dim, coords_l)]) missing_0 = xr.DataArray(coords_r, [(dim, coords_r)]) @@ -3727,7 +3755,7 @@ def test_binary_op_join_setting(self): expected = xr.DataArray([np.nan, 2, 4, np.nan], [(dim, [0, 1, 2, 3])]) assert_equal(actual, expected) - def test_combine_first(self): + def test_combine_first(self) -> None: ar0 = DataArray([[0, 0], [0, 0]], [("x", ["a", "b"]), ("y", [-1, 0])]) ar1 = DataArray([[1, 1], [1, 1]], [("x", ["b", "c"]), ("y", [0, 1])]) ar2 = DataArray([2], [("x", ["d"])]) @@ -3752,7 +3780,7 @@ def test_combine_first(self): ) assert_equal(actual, expected) - def test_sortby(self): + def test_sortby(self) -> None: da = DataArray( [[1, 2], [3, 4], [5, 6]], [("x", ["c", "b", "a"]), ("y", [1, 0])] ) @@ -3795,7 +3823,7 @@ def test_sortby(self): assert_equal(actual, expected) @requires_bottleneck - def test_rank(self): + def test_rank(self) -> None: # floats ar = DataArray([[3, 4, np.nan, 1]]) expect_0 = DataArray([[1, 1, np.nan, 1]]) @@ -3816,7 +3844,7 @@ def test_rank(self): @pytest.mark.parametrize("use_dask", [True, False]) @pytest.mark.parametrize("use_datetime", [True, False]) @pytest.mark.filterwarnings("ignore:overflow encountered in multiply") - def test_polyfit(self, use_dask, use_datetime): + def test_polyfit(self, use_dask, use_datetime) -> None: if use_dask and not has_dask: pytest.skip("requires dask") xcoord = xr.DataArray( @@ -3874,7 +3902,7 @@ def test_polyfit(self, use_dask, use_datetime): out = da.polyfit("x", 8, full=True) np.testing.assert_array_equal(out.polyfit_residuals.isnull(), [True, False]) - def test_pad_constant(self): + def test_pad_constant(self) -> None: ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) actual = ar.pad(dim_0=(1, 3)) expected = DataArray( @@ -3908,7 +3936,7 @@ def test_pad_constant(self): ) assert_identical(actual, expected) - def test_pad_coords(self): + def test_pad_coords(self) -> None: ar = DataArray( np.arange(3 * 4 * 5).reshape(3, 4, 5), [("x", np.arange(3)), ("y", np.arange(4)), ("z", np.arange(5))], @@ -3941,7 +3969,7 @@ def test_pad_coords(self): @pytest.mark.parametrize( "stat_length", (None, 3, (1, 3), {"dim_0": (2, 1), "dim_2": (4, 2)}) ) - def test_pad_stat_length(self, mode, stat_length): + def test_pad_stat_length(self, mode, stat_length) -> None: ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) actual = ar.pad(dim_0=(1, 3), dim_2=(2, 2), mode=mode, stat_length=stat_length) if isinstance(stat_length, dict): @@ -3960,7 +3988,7 @@ def test_pad_stat_length(self, mode, stat_length): @pytest.mark.parametrize( "end_values", (None, 3, (3, 5), {"dim_0": (2, 1), "dim_2": (4, 2)}) ) - def test_pad_linear_ramp(self, end_values): + def test_pad_linear_ramp(self, end_values) -> None: ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) actual = ar.pad( dim_0=(1, 3), dim_2=(2, 2), mode="linear_ramp", end_values=end_values @@ -3982,7 +4010,7 @@ def test_pad_linear_ramp(self, end_values): @pytest.mark.parametrize("mode", ("reflect", "symmetric")) @pytest.mark.parametrize("reflect_type", (None, "even", "odd")) - def test_pad_reflect(self, mode, reflect_type): + def test_pad_reflect(self, mode, reflect_type) -> None: ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) actual = ar.pad( @@ -4008,7 +4036,9 @@ def test_pad_reflect(self, mode, reflect_type): @pytest.mark.parametrize( "backend", ["numpy", pytest.param("dask", marks=[requires_dask])] ) - def test_query(self, backend, engine, parser): + def test_query( + self, backend, engine: QueryEngineOptions, parser: QueryParserOptions + ) -> None: """Test querying a dataset.""" # setup test data @@ -4063,7 +4093,7 @@ def test_query(self, backend, engine, parser): # test error handling with pytest.raises(ValueError): - aa.query("a > 5") # must be dict or kwargs + aa.query("a > 5") # type: ignore # must be dict or kwargs with pytest.raises(ValueError): aa.query(x=(a > 5)) # must be query string with pytest.raises(UndefinedVariableError): @@ -4071,7 +4101,7 @@ def test_query(self, backend, engine, parser): @requires_scipy @pytest.mark.parametrize("use_dask", [True, False]) - def test_curvefit(self, use_dask): + def test_curvefit(self, use_dask) -> None: if use_dask and not has_dask: pytest.skip("requires dask") @@ -4105,7 +4135,7 @@ def exp_decay(t, n0, tau=1): assert "a" in fit.param assert "x" not in fit.dims - def test_curvefit_helpers(self): + def test_curvefit_helpers(self) -> None: def exp_decay(t, n0, tau=1): return n0 * np.exp(-t / tau) @@ -4155,7 +4185,7 @@ def setup(self): ], ) class TestReduce1D(TestReduce): - def test_min(self, x, minindex, maxindex, nanindex): + def test_min(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs ) @@ -4181,7 +4211,7 @@ def test_min(self, x, minindex, maxindex, nanindex): assert_identical(result2, expected2) - def test_max(self, x, minindex, maxindex, nanindex): + def test_max(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs ) @@ -4210,7 +4240,7 @@ def test_max(self, x, minindex, maxindex, nanindex): @pytest.mark.filterwarnings( "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" ) - def test_argmin(self, x, minindex, maxindex, nanindex): + def test_argmin(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs ) @@ -4242,7 +4272,7 @@ def test_argmin(self, x, minindex, maxindex, nanindex): @pytest.mark.filterwarnings( "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" ) - def test_argmax(self, x, minindex, maxindex, nanindex): + def test_argmax(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs ) @@ -4272,7 +4302,7 @@ def test_argmax(self, x, minindex, maxindex, nanindex): assert_identical(result2, expected2) @pytest.mark.parametrize("use_dask", [True, False]) - def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): + def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask) -> None: if use_dask and not has_dask: pytest.skip("requires dask") if use_dask and x.dtype.kind == "M": @@ -4378,7 +4408,7 @@ def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): assert_identical(result7, expected7) @pytest.mark.parametrize("use_dask", [True, False]) - def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): + def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask) -> None: if use_dask and not has_dask: pytest.skip("requires dask") if use_dask and x.dtype.kind == "M": @@ -4486,7 +4516,7 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): @pytest.mark.filterwarnings( "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" ) - def test_argmin_dim(self, x, minindex, maxindex, nanindex): + def test_argmin_dim(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs ) @@ -4522,7 +4552,7 @@ def test_argmin_dim(self, x, minindex, maxindex, nanindex): @pytest.mark.filterwarnings( "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" ) - def test_argmax_dim(self, x, minindex, maxindex, nanindex): + def test_argmax_dim(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs ) @@ -4611,7 +4641,7 @@ def test_argmax_dim(self, x, minindex, maxindex, nanindex): ], ) class TestReduce2D(TestReduce): - def test_min(self, x, minindex, maxindex, nanindex): + def test_min(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["y", "x"], @@ -4620,10 +4650,10 @@ def test_min(self, x, minindex, maxindex, nanindex): ) minindex = [x if not np.isnan(x) else 0 for x in minindex] - expected0 = [ + expected0list = [ ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) ] - expected0 = xr.concat(expected0, dim="y") + expected0 = xr.concat(expected0list, dim="y") result0 = ar.min(dim="x", keep_attrs=True) assert_identical(result0, expected0) @@ -4640,17 +4670,17 @@ def test_min(self, x, minindex, maxindex, nanindex): x if y is None or ar.dtype.kind == "O" else y for x, y in zip(minindex, nanindex) ] - expected2 = [ + expected2list = [ ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) ] - expected2 = xr.concat(expected2, dim="y") + expected2 = xr.concat(expected2list, dim="y") expected2.attrs = {} result3 = ar.min(dim="x", skipna=False) assert_identical(result3, expected2) - def test_max(self, x, minindex, maxindex, nanindex): + def test_max(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["y", "x"], @@ -4659,10 +4689,10 @@ def test_max(self, x, minindex, maxindex, nanindex): ) maxindex = [x if not np.isnan(x) else 0 for x in maxindex] - expected0 = [ + expected0list = [ ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) ] - expected0 = xr.concat(expected0, dim="y") + expected0 = xr.concat(expected0list, dim="y") result0 = ar.max(dim="x", keep_attrs=True) assert_identical(result0, expected0) @@ -4679,36 +4709,36 @@ def test_max(self, x, minindex, maxindex, nanindex): x if y is None or ar.dtype.kind == "O" else y for x, y in zip(maxindex, nanindex) ] - expected2 = [ + expected2list = [ ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) ] - expected2 = xr.concat(expected2, dim="y") + expected2 = xr.concat(expected2list, dim="y") expected2.attrs = {} result3 = ar.max(dim="x", skipna=False) assert_identical(result3, expected2) - def test_argmin(self, x, minindex, maxindex, nanindex): + def test_argmin(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["y", "x"], coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, attrs=self.attrs, ) - indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) - indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) + indarrnp = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) + indarr = xr.DataArray(indarrnp, dims=ar.dims, coords=ar.coords) if np.isnan(minindex).any(): with pytest.raises(ValueError): ar.argmin(dim="x") return - expected0 = [ + expected0list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) ] - expected0 = xr.concat(expected0, dim="y") + expected0 = xr.concat(expected0list, dim="y") result0 = ar.argmin(dim="x") assert_identical(result0, expected0) @@ -4725,37 +4755,37 @@ def test_argmin(self, x, minindex, maxindex, nanindex): x if y is None or ar.dtype.kind == "O" else y for x, y in zip(minindex, nanindex) ] - expected2 = [ + expected2list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) ] - expected2 = xr.concat(expected2, dim="y") + expected2 = xr.concat(expected2list, dim="y") expected2.attrs = {} result3 = ar.argmin(dim="x", skipna=False) assert_identical(result3, expected2) - def test_argmax(self, x, minindex, maxindex, nanindex): + def test_argmax(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["y", "x"], coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, attrs=self.attrs, ) - indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) - indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) + indarr_np = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) + indarr = xr.DataArray(indarr_np, dims=ar.dims, coords=ar.coords) if np.isnan(maxindex).any(): with pytest.raises(ValueError): ar.argmax(dim="x") return - expected0 = [ + expected0list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) ] - expected0 = xr.concat(expected0, dim="y") + expected0 = xr.concat(expected0list, dim="y") result0 = ar.argmax(dim="x") assert_identical(result0, expected0) @@ -4772,11 +4802,11 @@ def test_argmax(self, x, minindex, maxindex, nanindex): x if y is None or ar.dtype.kind == "O" else y for x, y in zip(maxindex, nanindex) ] - expected2 = [ + expected2list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) ] - expected2 = xr.concat(expected2, dim="y") + expected2 = xr.concat(expected2list, dim="y") expected2.attrs = {} result3 = ar.argmax(dim="x", skipna=False) @@ -4784,7 +4814,7 @@ def test_argmax(self, x, minindex, maxindex, nanindex): assert_identical(result3, expected2) @pytest.mark.parametrize("use_dask", [True, False]) - def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): + def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask) -> None: if use_dask and not has_dask: pytest.skip("requires dask") if use_dask and x.dtype.kind == "M": @@ -4830,11 +4860,11 @@ def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): minindex0 = [x if not np.isnan(x) else 0 for x in minindex] nan_mult_0 = np.array([np.NaN if x else 1 for x in hasna])[:, None] - expected0 = [ + expected0list = [ (coordarr1 * nan_mult_0).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex0) ] - expected0 = xr.concat(expected0, dim="y") + expected0 = xr.concat(expected0list, dim="y") expected0.name = "x" # Default fill value (NaN) @@ -4859,11 +4889,11 @@ def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): x if y is None or ar0.dtype.kind == "O" else y for x, y in zip(minindex0, nanindex) ] - expected3 = [ + expected3list = [ coordarr0.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex3) ] - expected3 = xr.concat(expected3, dim="y") + expected3 = xr.concat(expected3list, dim="y") expected3.name = "x" expected3.attrs = {} @@ -4878,11 +4908,11 @@ def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): # Float fill_value nan_mult_5 = np.array([-1.1 if x else 1 for x in hasna])[:, None] - expected5 = [ + expected5list = [ (coordarr1 * nan_mult_5).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex0) ] - expected5 = xr.concat(expected5, dim="y") + expected5 = xr.concat(expected5list, dim="y") expected5.name = "x" with raise_if_dask_computes(max_computes=max_computes): @@ -4891,11 +4921,11 @@ def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): # Integer fill_value nan_mult_6 = np.array([-1 if x else 1 for x in hasna])[:, None] - expected6 = [ + expected6list = [ (coordarr1 * nan_mult_6).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex0) ] - expected6 = xr.concat(expected6, dim="y") + expected6 = xr.concat(expected6list, dim="y") expected6.name = "x" with raise_if_dask_computes(max_computes=max_computes): @@ -4904,11 +4934,11 @@ def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): # Complex fill_value nan_mult_7 = np.array([-5j if x else 1 for x in hasna])[:, None] - expected7 = [ + expected7list = [ (coordarr1 * nan_mult_7).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex0) ] - expected7 = xr.concat(expected7, dim="y") + expected7 = xr.concat(expected7list, dim="y") expected7.name = "x" with raise_if_dask_computes(max_computes=max_computes): @@ -4916,7 +4946,7 @@ def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): assert_identical(result7, expected7) @pytest.mark.parametrize("use_dask", [True, False]) - def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): + def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask) -> None: if use_dask and not has_dask: pytest.skip("requires dask") if use_dask and x.dtype.kind == "M": @@ -4963,11 +4993,11 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): maxindex0 = [x if not np.isnan(x) else 0 for x in maxindex] nan_mult_0 = np.array([np.NaN if x else 1 for x in hasna])[:, None] - expected0 = [ + expected0list = [ (coordarr1 * nan_mult_0).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex0) ] - expected0 = xr.concat(expected0, dim="y") + expected0 = xr.concat(expected0list, dim="y") expected0.name = "x" # Default fill value (NaN) @@ -4992,11 +5022,11 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): x if y is None or ar0.dtype.kind == "O" else y for x, y in zip(maxindex0, nanindex) ] - expected3 = [ + expected3list = [ coordarr0.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex3) ] - expected3 = xr.concat(expected3, dim="y") + expected3 = xr.concat(expected3list, dim="y") expected3.name = "x" expected3.attrs = {} @@ -5011,11 +5041,11 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): # Float fill_value nan_mult_5 = np.array([-1.1 if x else 1 for x in hasna])[:, None] - expected5 = [ + expected5list = [ (coordarr1 * nan_mult_5).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex0) ] - expected5 = xr.concat(expected5, dim="y") + expected5 = xr.concat(expected5list, dim="y") expected5.name = "x" with raise_if_dask_computes(max_computes=max_computes): @@ -5024,11 +5054,11 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): # Integer fill_value nan_mult_6 = np.array([-1 if x else 1 for x in hasna])[:, None] - expected6 = [ + expected6list = [ (coordarr1 * nan_mult_6).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex0) ] - expected6 = xr.concat(expected6, dim="y") + expected6 = xr.concat(expected6list, dim="y") expected6.name = "x" with raise_if_dask_computes(max_computes=max_computes): @@ -5037,11 +5067,11 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): # Complex fill_value nan_mult_7 = np.array([-5j if x else 1 for x in hasna])[:, None] - expected7 = [ + expected7list = [ (coordarr1 * nan_mult_7).isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex0) ] - expected7 = xr.concat(expected7, dim="y") + expected7 = xr.concat(expected7list, dim="y") expected7.name = "x" with raise_if_dask_computes(max_computes=max_computes): @@ -5051,26 +5081,26 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): @pytest.mark.filterwarnings( "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" ) - def test_argmin_dim(self, x, minindex, maxindex, nanindex): + def test_argmin_dim(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["y", "x"], coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, attrs=self.attrs, ) - indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) - indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) + indarrnp = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) + indarr = xr.DataArray(indarrnp, dims=ar.dims, coords=ar.coords) if np.isnan(minindex).any(): with pytest.raises(ValueError): ar.argmin(dim="x") return - expected0 = [ + expected0list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) ] - expected0 = {"x": xr.concat(expected0, dim="y")} + expected0 = {"x": xr.concat(expected0list, dim="y")} result0 = ar.argmin(dim=["x"]) for key in expected0: @@ -5086,11 +5116,11 @@ def test_argmin_dim(self, x, minindex, maxindex, nanindex): x if y is None or ar.dtype.kind == "O" else y for x, y in zip(minindex, nanindex) ] - expected2 = [ + expected2list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) ] - expected2 = {"x": xr.concat(expected2, dim="y")} + expected2 = {"x": xr.concat(expected2list, dim="y")} expected2["x"].attrs = {} result2 = ar.argmin(dim=["x"], skipna=False) @@ -5099,7 +5129,8 @@ def test_argmin_dim(self, x, minindex, maxindex, nanindex): assert_identical(result2[key], expected2[key]) result3 = ar.argmin(...) - min_xind = ar.isel(expected0).argmin() + # TODO: remove cast once argmin typing is overloaded + min_xind = cast(DataArray, ar.isel(expected0).argmin()) expected3 = { "y": DataArray(min_xind), "x": DataArray(minindex[min_xind.item()]), @@ -5111,26 +5142,26 @@ def test_argmin_dim(self, x, minindex, maxindex, nanindex): @pytest.mark.filterwarnings( "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" ) - def test_argmax_dim(self, x, minindex, maxindex, nanindex): + def test_argmax_dim(self, x, minindex, maxindex, nanindex) -> None: ar = xr.DataArray( x, dims=["y", "x"], coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, attrs=self.attrs, ) - indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) - indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) + indarrnp = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) + indarr = xr.DataArray(indarrnp, dims=ar.dims, coords=ar.coords) if np.isnan(maxindex).any(): with pytest.raises(ValueError): ar.argmax(dim="x") return - expected0 = [ + expected0list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) ] - expected0 = {"x": xr.concat(expected0, dim="y")} + expected0 = {"x": xr.concat(expected0list, dim="y")} result0 = ar.argmax(dim=["x"]) for key in expected0: @@ -5146,11 +5177,11 @@ def test_argmax_dim(self, x, minindex, maxindex, nanindex): x if y is None or ar.dtype.kind == "O" else y for x, y in zip(maxindex, nanindex) ] - expected2 = [ + expected2list = [ indarr.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) ] - expected2 = {"x": xr.concat(expected2, dim="y")} + expected2 = {"x": xr.concat(expected2list, dim="y")} expected2["x"].attrs = {} result2 = ar.argmax(dim=["x"], skipna=False) @@ -5159,7 +5190,8 @@ def test_argmax_dim(self, x, minindex, maxindex, nanindex): assert_identical(result2[key], expected2[key]) result3 = ar.argmax(...) - max_xind = ar.isel(expected0).argmax() + # TODO: remove cast once argmax typing is overloaded + max_xind = cast(DataArray, ar.isel(expected0).argmax()) expected3 = { "y": DataArray(max_xind), "x": DataArray(maxindex[max_xind.item()]), @@ -5797,7 +5829,7 @@ def test_argmax_dim( class TestReduceND(TestReduce): @pytest.mark.parametrize("op", ["idxmin", "idxmax"]) @pytest.mark.parametrize("ndim", [3, 5]) - def test_idxminmax_dask(self, op, ndim): + def test_idxminmax_dask(self, op, ndim) -> None: if not has_dask: pytest.skip("requires dask") @@ -5842,7 +5874,7 @@ def da(request, backend): @pytest.mark.parametrize("da", ("repeating_ints",), indirect=True) -def test_isin(da): +def test_isin(da) -> None: expected = DataArray( np.asarray([[0, 0, 0], [1, 0, 0]]), dims=list("yx"), @@ -5862,7 +5894,7 @@ def test_isin(da): @pytest.mark.parametrize("da", (1, 2), indirect=True) -def test_rolling_iter(da): +def test_rolling_iter(da) -> None: rolling_obj = da.rolling(time=7) rolling_obj_mean = rolling_obj.mean() @@ -5886,7 +5918,7 @@ def test_rolling_iter(da): @pytest.mark.parametrize("da", (1,), indirect=True) -def test_rolling_repr(da): +def test_rolling_repr(da) -> None: rolling_obj = da.rolling(time=7) assert repr(rolling_obj) == "DataArrayRolling [time->7]" rolling_obj = da.rolling(time=7, center=True) @@ -5896,7 +5928,7 @@ def test_rolling_repr(da): @requires_dask -def test_repeated_rolling_rechunks(): +def test_repeated_rolling_rechunks() -> None: # regression test for GH3277, GH2514 dat = DataArray(np.random.rand(7653, 300), dims=("day", "item")) @@ -5904,14 +5936,14 @@ def test_repeated_rolling_rechunks(): dat_chunk.rolling(day=10).mean().rolling(day=250).std() -def test_rolling_doc(da): +def test_rolling_doc(da) -> None: rolling_obj = da.rolling(time=7) # argument substitution worked assert "`mean`" in rolling_obj.mean.__doc__ -def test_rolling_properties(da): +def test_rolling_properties(da) -> None: rolling_obj = da.rolling(time=4) assert rolling_obj.obj.get_axis_num("time") == 1 @@ -5928,7 +5960,7 @@ def test_rolling_properties(da): @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) -def test_rolling_wrapped_bottleneck(da, name, center, min_periods): +def test_rolling_wrapped_bottleneck(da, name, center, min_periods) -> None: bn = pytest.importorskip("bottleneck", minversion="1.1") # Test all bottleneck functions @@ -5956,7 +5988,7 @@ def test_rolling_wrapped_bottleneck(da, name, center, min_periods): @pytest.mark.parametrize("min_periods", (1, None)) @pytest.mark.parametrize("window", (7, 8)) @pytest.mark.parametrize("backend", ["dask"], indirect=True) -def test_rolling_wrapped_dask(da, name, center, min_periods, window): +def test_rolling_wrapped_dask(da, name, center, min_periods, window) -> None: # dask version rolling_obj = da.rolling(time=window, min_periods=min_periods, center=center) actual = getattr(rolling_obj, name)().load() @@ -5980,7 +6012,7 @@ def test_rolling_wrapped_dask(da, name, center, min_periods, window): @pytest.mark.parametrize("center", (True, None)) -def test_rolling_wrapped_dask_nochunk(center): +def test_rolling_wrapped_dask_nochunk(center) -> None: # GH:2113 pytest.importorskip("dask.array") @@ -5995,7 +6027,7 @@ def test_rolling_wrapped_dask_nochunk(center): @pytest.mark.parametrize("center", (True, False)) @pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) -def test_rolling_pandas_compat(center, window, min_periods): +def test_rolling_pandas_compat(center, window, min_periods) -> None: s = pd.Series(np.arange(10)) da = DataArray.from_series(s) @@ -6016,7 +6048,7 @@ def test_rolling_pandas_compat(center, window, min_periods): @pytest.mark.parametrize("center", (True, False)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) -def test_rolling_construct(center, window): +def test_rolling_construct(center, window) -> None: s = pd.Series(np.arange(10)) da = DataArray.from_series(s) @@ -6045,7 +6077,7 @@ def test_rolling_construct(center, window): @pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) @pytest.mark.parametrize("name", ("sum", "mean", "std", "max")) -def test_rolling_reduce(da, center, min_periods, window, name): +def test_rolling_reduce(da, center, min_periods, window, name) -> None: if min_periods is not None and window < min_periods: min_periods = window @@ -6066,7 +6098,7 @@ def test_rolling_reduce(da, center, min_periods, window, name): @pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) @pytest.mark.parametrize("window", (1, 2, 3, 4)) @pytest.mark.parametrize("name", ("sum", "max")) -def test_rolling_reduce_nonnumeric(center, min_periods, window, name): +def test_rolling_reduce_nonnumeric(center, min_periods, window, name) -> None: da = DataArray( [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time" ).isnull() @@ -6083,10 +6115,10 @@ def test_rolling_reduce_nonnumeric(center, min_periods, window, name): assert actual.dims == expected.dims -def test_rolling_count_correct(): +def test_rolling_count_correct() -> None: da = DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time") - kwargs = [ + kwargs: list[dict[str, Any]] = [ {"time": 11, "min_periods": 1}, {"time": 11, "min_periods": None}, {"time": 7, "min_periods": 2}, @@ -6124,7 +6156,7 @@ def test_rolling_count_correct(): @pytest.mark.parametrize("center", (True, False)) @pytest.mark.parametrize("min_periods", (None, 1)) @pytest.mark.parametrize("name", ("sum", "mean", "max")) -def test_ndrolling_reduce(da, center, min_periods, name): +def test_ndrolling_reduce(da, center, min_periods, name) -> None: rolling_obj = da.rolling(time=3, x=2, center=center, min_periods=min_periods) actual = getattr(rolling_obj, name)() @@ -6151,7 +6183,7 @@ def test_ndrolling_reduce(da, center, min_periods, name): @pytest.mark.parametrize("center", (True, False, (True, False))) @pytest.mark.parametrize("fill_value", (np.nan, 0.0)) -def test_ndrolling_construct(center, fill_value): +def test_ndrolling_construct(center, fill_value) -> None: da = DataArray( np.arange(5 * 6 * 7).reshape(5, 6, 7).astype(float), dims=["x", "y", "z"], @@ -6180,7 +6212,7 @@ def test_ndrolling_construct(center, fill_value): ("count", ()), ], ) -def test_rolling_keep_attrs(funcname, argument): +def test_rolling_keep_attrs(funcname, argument) -> None: attrs_da = {"da_attr": "test"} data = np.linspace(10, 15, 100) @@ -6223,17 +6255,17 @@ def test_rolling_keep_attrs(funcname, argument): assert result.name == "name" -def test_raise_no_warning_for_nan_in_binary_ops(): +def test_raise_no_warning_for_nan_in_binary_ops() -> None: with assert_no_warnings(): xr.DataArray([1, 2, np.NaN]) > 0 @pytest.mark.filterwarnings("error") -def test_no_warning_for_all_nan(): +def test_no_warning_for_all_nan() -> None: _ = xr.DataArray([np.NaN, np.NaN]).mean() -def test_name_in_masking(): +def test_name_in_masking() -> None: name = "RingoStarr" da = xr.DataArray(range(10), coords=[("x", range(10))], name=name) assert da.where(da > 5).name == name @@ -6244,12 +6276,12 @@ def test_name_in_masking(): class TestIrisConversion: @requires_iris - def test_to_and_from_iris(self): + def test_to_and_from_iris(self) -> None: import cf_units # iris requirement import iris # to iris - coord_dict = {} + coord_dict: dict[Hashable, Any] = {} coord_dict["distance"] = ("distance", [-2, 2], {"units": "meters"}) coord_dict["time"] = ("time", pd.date_range("2000-01-01", periods=3)) coord_dict["height"] = 10 @@ -6315,12 +6347,12 @@ def test_to_and_from_iris(self): @requires_iris @requires_dask - def test_to_and_from_iris_dask(self): + def test_to_and_from_iris_dask(self) -> None: import cf_units # iris requirement import dask.array as da import iris - coord_dict = {} + coord_dict: dict[Hashable, Any] = {} coord_dict["distance"] = ("distance", [-2, 2], {"units": "meters"}) coord_dict["time"] = ("time", pd.date_range("2000-01-01", periods=3)) coord_dict["height"] = 10 @@ -6417,15 +6449,14 @@ def test_to_and_from_iris_dask(self): (None, None, None, None, {}), ], ) - def test_da_name_from_cube(self, std_name, long_name, var_name, name, attrs): + def test_da_name_from_cube( + self, std_name, long_name, var_name, name, attrs + ) -> None: from iris.cube import Cube - data = [] - cube = Cube( - data, var_name=var_name, standard_name=std_name, long_name=long_name - ) + cube = Cube([], var_name=var_name, standard_name=std_name, long_name=long_name) result = xr.DataArray.from_iris(cube) - expected = xr.DataArray(data, name=name, attrs=attrs) + expected = xr.DataArray([], name=name, attrs=attrs) xr.testing.assert_identical(result, expected) @requires_iris @@ -6450,7 +6481,9 @@ def test_da_name_from_cube(self, std_name, long_name, var_name, name, attrs): (None, None, None, "unknown", {}), ], ) - def test_da_coord_name_from_cube(self, std_name, long_name, var_name, name, attrs): + def test_da_coord_name_from_cube( + self, std_name, long_name, var_name, name, attrs + ) -> None: from iris.coords import DimCoord from iris.cube import Cube @@ -6464,7 +6497,7 @@ def test_da_coord_name_from_cube(self, std_name, long_name, var_name, name, attr xr.testing.assert_identical(result, expected) @requires_iris - def test_prevent_duplicate_coord_names(self): + def test_prevent_duplicate_coord_names(self) -> None: from iris.coords import DimCoord from iris.cube import Cube @@ -6486,7 +6519,7 @@ def test_prevent_duplicate_coord_names(self): "coord_values", [["IA", "IL", "IN"], [0, 2, 1]], # non-numeric values # non-monotonic values ) - def test_fallback_to_iris_AuxCoord(self, coord_values): + def test_fallback_to_iris_AuxCoord(self, coord_values) -> None: from iris.coords import AuxCoord from iris.cube import Cube @@ -6506,7 +6539,7 @@ def test_fallback_to_iris_AuxCoord(self, coord_values): ) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) @pytest.mark.parametrize("func", ["mean", "sum"]) -def test_rolling_exp_runs(da, dim, window_type, window, func): +def test_rolling_exp_runs(da, dim, window_type, window, func) -> None: import numbagg if ( @@ -6528,7 +6561,7 @@ def test_rolling_exp_runs(da, dim, window_type, window, func): "window_type, window", [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]] ) @pytest.mark.parametrize("backend", ["numpy"], indirect=True) -def test_rolling_exp_mean_pandas(da, dim, window_type, window): +def test_rolling_exp_mean_pandas(da, dim, window_type, window) -> None: da = da.isel(a=0).where(lambda x: x > 0.2) result = da.rolling_exp(window_type=window_type, **{dim: window}).mean() @@ -6548,7 +6581,7 @@ def test_rolling_exp_mean_pandas(da, dim, window_type, window): @requires_numbagg @pytest.mark.parametrize("backend", ["numpy"], indirect=True) @pytest.mark.parametrize("func", ["mean", "sum"]) -def test_rolling_exp_keep_attrs(da, func): +def test_rolling_exp_keep_attrs(da, func) -> None: import numbagg if ( @@ -6591,13 +6624,13 @@ def test_rolling_exp_keep_attrs(da, func): da.rolling_exp(time=10, keep_attrs=True) -def test_no_dict(): +def test_no_dict() -> None: d = DataArray() with pytest.raises(AttributeError): d.__dict__ -def test_subclass_slots(): +def test_subclass_slots() -> None: """Test that DataArray subclasses must explicitly define ``__slots__``. .. note:: @@ -6612,7 +6645,7 @@ class MyArray(DataArray): assert str(e.value) == "MyArray must explicitly define __slots__" -def test_weakref(): +def test_weakref() -> None: """Classes with __slots__ are incompatible with the weakref module unless they explicitly state __weakref__ among their slots """ @@ -6623,7 +6656,7 @@ def test_weakref(): assert r() is a -def test_delete_coords(): +def test_delete_coords() -> None: """Make sure that deleting a coordinate doesn't corrupt the DataArray. See issue #3899. @@ -6648,13 +6681,13 @@ def test_delete_coords(): assert set(a1.coords.keys()) == {"x"} -def test_deepcopy_obj_array(): +def test_deepcopy_obj_array() -> None: x0 = DataArray(np.array([object()])) x1 = deepcopy(x0) assert x0.values[0] is not x1.values[0] -def test_clip(da): +def test_clip(da) -> None: with raise_if_dask_computes(): result = da.clip(min=0.5) assert result.min(...) >= 0.5 @@ -6688,7 +6721,7 @@ def test_clip(da): class TestDropDuplicates: @pytest.mark.parametrize("keep", ["first", "last", False]) - def test_drop_duplicates_1d(self, keep): + def test_drop_duplicates_1d(self, keep) -> None: da = xr.DataArray( [0, 5, 6, 7], dims="time", coords={"time": [0, 0, 1, 2]}, name="test" ) @@ -6710,7 +6743,7 @@ def test_drop_duplicates_1d(self, keep): with pytest.raises(ValueError, match="['space'] not found"): da.drop_duplicates("space", keep=keep) - def test_drop_duplicates_2d(self): + def test_drop_duplicates_2d(self) -> None: da = xr.DataArray( [[0, 5, 6, 7], [2, 1, 3, 4]], dims=["space", "time"], @@ -6734,7 +6767,7 @@ def test_drop_duplicates_2d(self): class TestNumpyCoercion: # TODO once flexible indexes refactor complete also test coercion of dimension coords - def test_from_numpy(self): + def test_from_numpy(self) -> None: da = xr.DataArray([1, 2, 3], dims="x", coords={"lat": ("x", [4, 5, 6])}) assert_identical(da.as_numpy(), da) @@ -6742,7 +6775,7 @@ def test_from_numpy(self): np.testing.assert_equal(da["lat"].to_numpy(), np.array([4, 5, 6])) @requires_dask - def test_from_dask(self): + def test_from_dask(self) -> None: da = xr.DataArray([1, 2, 3], dims="x", coords={"lat": ("x", [4, 5, 6])}) da_chunked = da.chunk(1) @@ -6751,7 +6784,7 @@ def test_from_dask(self): np.testing.assert_equal(da["lat"].to_numpy(), np.array([4, 5, 6])) @requires_pint - def test_from_pint(self): + def test_from_pint(self) -> None: from pint import Quantity arr = np.array([1, 2, 3]) @@ -6767,7 +6800,7 @@ def test_from_pint(self): np.testing.assert_equal(da["lat"].to_numpy(), arr + 3) @requires_sparse - def test_from_sparse(self): + def test_from_sparse(self) -> None: import sparse arr = np.diagflat([1, 2, 3]) @@ -6783,7 +6816,7 @@ def test_from_sparse(self): np.testing.assert_equal(da.to_numpy(), arr) @requires_cupy - def test_from_cupy(self): + def test_from_cupy(self) -> None: import cupy as cp arr = np.array([1, 2, 3]) @@ -6797,7 +6830,7 @@ def test_from_cupy(self): @requires_dask @requires_pint - def test_from_pint_wrapping_dask(self): + def test_from_pint_wrapping_dask(self) -> None: import dask from pint import Quantity @@ -6818,13 +6851,13 @@ def test_from_pint_wrapping_dask(self): class TestStackEllipsis: # https://github.com/pydata/xarray/issues/6051 - def test_result_as_expected(self): + def test_result_as_expected(self) -> None: da = DataArray([[1, 2], [1, 2]], dims=("x", "y")) result = da.stack(flat=[...]) expected = da.stack(flat=da.dims) assert_identical(result, expected) - def test_error_on_ellipsis_without_list(self): + def test_error_on_ellipsis_without_list(self) -> None: da = DataArray([[1, 2], [1, 2]], dims=("x", "y")) with pytest.raises(ValueError): - da.stack(flat=...) + da.stack(flat=...) # type: ignore diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 5f368375fc0..950f15e91df 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -76,6 +76,8 @@ def create_append_test_data(seed=None): time2 = pd.date_range("2000-02-01", periods=nt2) string_var = np.array(["ae", "bc", "df"], dtype=object) string_var_to_append = np.array(["asdf", "asdfg"], dtype=object) + string_var_fixed_length = np.array(["aa", "bb", "cc"], dtype="|S2") + string_var_fixed_length_to_append = np.array(["dd", "ee"], dtype="|S2") unicode_var = ["áó", "áó", "áó"] datetime_var = np.array( ["2019-01-01", "2019-01-02", "2019-01-03"], dtype="datetime64[s]" @@ -94,6 +96,9 @@ def create_append_test_data(seed=None): dims=["lat", "lon", "time"], ), "string_var": xr.DataArray(string_var, coords=[time1], dims=["time"]), + "string_var_fixed_length": xr.DataArray( + string_var_fixed_length, coords=[time1], dims=["time"] + ), "unicode_var": xr.DataArray( unicode_var, coords=[time1], dims=["time"] ).astype(np.unicode_), @@ -112,6 +117,9 @@ def create_append_test_data(seed=None): "string_var": xr.DataArray( string_var_to_append, coords=[time2], dims=["time"] ), + "string_var_fixed_length": xr.DataArray( + string_var_fixed_length_to_append, coords=[time2], dims=["time"] + ), "unicode_var": xr.DataArray( unicode_var[:nt2], coords=[time2], dims=["time"] ).astype(np.unicode_), @@ -137,6 +145,33 @@ def create_append_test_data(seed=None): return ds, ds_to_append, ds_with_new_var +def create_append_string_length_mismatch_test_data(dtype): + def make_datasets(data, data_to_append): + ds = xr.Dataset( + {"temperature": (["time"], data)}, + coords={"time": [0, 1, 2]}, + ) + ds_to_append = xr.Dataset( + {"temperature": (["time"], data_to_append)}, coords={"time": [0, 1, 2]} + ) + assert all(objp.data.flags.writeable for objp in ds.variables.values()) + assert all( + objp.data.flags.writeable for objp in ds_to_append.variables.values() + ) + return ds, ds_to_append + + u2_strings = ["ab", "cd", "ef"] + u5_strings = ["abc", "def", "ghijk"] + + s2_strings = np.array(["aa", "bb", "cc"], dtype="|S2") + s3_strings = np.array(["aaa", "bbb", "ccc"], dtype="|S3") + + if dtype == "U": + return make_datasets(u2_strings, u5_strings) + elif dtype == "S": + return make_datasets(s2_strings, s3_strings) + + def create_test_multiindex(): mindex = pd.MultiIndex.from_product( [["a", "b"], [1, 2]], names=("level_1", "level_2") @@ -921,6 +956,9 @@ def test_chunk(self): expected_chunks = {"dim1": (8,), "dim2": (9,), "dim3": (10,)} assert reblocked.chunks == expected_chunks + # test kwargs form of chunks + assert data.chunk(**expected_chunks).chunks == expected_chunks + def get_dask_names(ds): return {k: v.data.name for k, v in ds.items()} @@ -947,7 +985,7 @@ def get_dask_names(ds): new_dask_names = get_dask_names(reblocked) assert reblocked.chunks == expected_chunks assert_identical(reblocked, data) - # recuhnking with same chunk sizes should not change names + # rechunking with same chunk sizes should not change names for k, v in new_dask_names.items(): assert v == orig_dask_names[k] @@ -1172,6 +1210,25 @@ def test_isel_fancy(self): assert_array_equal(actual["var2"], expected_var2) assert_array_equal(actual["var3"], expected_var3) + # test that drop works + ds = xr.Dataset({"a": (("x",), [1, 2, 3])}, coords={"b": (("x",), [5, 6, 7])}) + + actual = ds.isel({"x": 1}, drop=False) + expected = xr.Dataset({"a": 2}, coords={"b": 6}) + assert_identical(actual, expected) + + actual = ds.isel({"x": 1}, drop=True) + expected = xr.Dataset({"a": 2}) + assert_identical(actual, expected) + + actual = ds.isel({"x": DataArray(1)}, drop=False) + expected = xr.Dataset({"a": 2}, coords={"b": 6}) + assert_identical(actual, expected) + + actual = ds.isel({"x": DataArray(1)}, drop=True) + expected = xr.Dataset({"a": 2}) + assert_identical(actual, expected) + def test_isel_dataarray(self): """Test for indexing by DataArray""" data = create_test_data() @@ -2328,6 +2385,18 @@ def test_broadcast_misaligned(self): assert_identical(expected_x2, x2) assert_identical(expected_y2, y2) + def test_broadcast_multi_index(self): + # GH6430 + ds = Dataset( + {"foo": (("x", "y", "z"), np.ones((3, 4, 2)))}, + {"x": ["a", "b", "c"], "y": [1, 2, 3, 4]}, + ) + stacked = ds.stack(space=["x", "y"]) + broadcasted, _ = broadcast(stacked, stacked.space) + + assert broadcasted.xindexes["x"] is broadcasted.xindexes["space"] + assert broadcasted.xindexes["y"] is broadcasted.xindexes["space"] + def test_variable_indexing(self): data = create_test_data() v = data["var1"] @@ -2384,11 +2453,10 @@ def test_drop_variables(self): def test_drop_multiindex_level(self): data = create_test_multiindex() - - with pytest.raises( - ValueError, match=r"cannot remove coordinate.*corrupt.*index " - ): - data.drop_vars("level_1") + expected = data.drop_vars(["x", "level_1", "level_2"]) + with pytest.warns(DeprecationWarning): + actual = data.drop_vars("level_1") + assert_identical(expected, actual) def test_drop_index_labels(self): data = Dataset({"A": (["x", "y"], np.random.randn(2, 3)), "x": ["a", "b"]}) @@ -4545,8 +4613,11 @@ def test_where_other(self): actual = ds.where(lambda x: x > 1, -1) assert_equal(expected, actual) - with pytest.raises(ValueError, match=r"cannot set"): - ds.where(ds > 1, other=0, drop=True) + actual = ds.where(ds > 1, other=-1, drop=True) + expected_nodrop = ds.where(ds > 1, -1) + _, expected = xr.align(actual, expected_nodrop, join="left") + assert_equal(actual, expected) + assert actual.a.dtype == int with pytest.raises(ValueError, match=r"cannot align .* are not equal"): ds.where(ds > 1, ds.isel(x=slice(3))) @@ -5488,7 +5559,7 @@ def test_binary_op_join_setting(self): actual = ds1 + ds2 assert_equal(actual, expected) - def test_full_like(self): + def test_full_like(self) -> None: # For more thorough tests, see test_variable.py # Note: testing data_vars with mismatched dtypes ds = Dataset( @@ -5501,8 +5572,9 @@ def test_full_like(self): actual = full_like(ds, 2) expected = ds.copy(deep=True) - expected["d1"].values = [2, 2, 2] - expected["d2"].values = [2.0, 2.0, 2.0] + # https://github.com/python/mypy/issues/3004 + expected["d1"].values = [2, 2, 2] # type: ignore + expected["d2"].values = [2.0, 2.0, 2.0] # type: ignore assert expected["d1"].dtype == int assert expected["d2"].dtype == float assert_identical(expected, actual) @@ -5510,8 +5582,8 @@ def test_full_like(self): # override dtype actual = full_like(ds, fill_value=True, dtype=bool) expected = ds.copy(deep=True) - expected["d1"].values = [True, True, True] - expected["d2"].values = [True, True, True] + expected["d1"].values = [True, True, True] # type: ignore + expected["d2"].values = [True, True, True] # type: ignore assert expected["d1"].dtype == bool assert expected["d2"].dtype == bool assert_identical(expected, actual) @@ -5717,7 +5789,7 @@ def test_ipython_key_completion(self): ds.data_vars[item] # should not raise assert sorted(actual) == sorted(expected) - def test_polyfit_output(self): + def test_polyfit_output(self) -> None: ds = create_test_data(seed=1) out = ds.polyfit("dim2", 2, full=False) @@ -5730,7 +5802,7 @@ def test_polyfit_output(self): out = ds.polyfit("time", 2) assert len(out.data_vars) == 0 - def test_polyfit_warnings(self): + def test_polyfit_warnings(self) -> None: ds = create_test_data(seed=1) with warnings.catch_warnings(record=True) as ws: diff --git a/xarray/tests/test_distributed.py b/xarray/tests/test_distributed.py index 773733b7b89..cde24c101ea 100644 --- a/xarray/tests/test_distributed.py +++ b/xarray/tests/test_distributed.py @@ -2,6 +2,8 @@ import pickle import numpy as np +from typing import Any + import pytest from packaging.version import Version @@ -10,7 +12,7 @@ from dask.distributed import Client, Lock from distributed.client import futures_of -from distributed.utils_test import cluster, gen_cluster, loop +from distributed.utils_test import cluster, gen_cluster, loop, cleanup # noqa: F401 import xarray as xr from xarray.backends.locks import HDF5_LOCK, CombinedLock @@ -156,7 +158,7 @@ def test_dask_distributed_zarr_integration_test(loop, consolidated, compute) -> if consolidated: pytest.importorskip("zarr", minversion="2.2.1.dev2") write_kwargs = {"consolidated": True} - read_kwargs = {"backend_kwargs": {"consolidated": True}} + read_kwargs: dict[str, Any] = {"backend_kwargs": {"consolidated": True}} else: write_kwargs = read_kwargs = {} # type: ignore chunks = {"dim1": 4, "dim2": 3, "dim3": 5} diff --git a/xarray/tests/test_duck_array_ops.py b/xarray/tests/test_duck_array_ops.py index c329bc50c56..392f1b91914 100644 --- a/xarray/tests/test_duck_array_ops.py +++ b/xarray/tests/test_duck_array_ops.py @@ -675,39 +675,68 @@ def test_multiple_dims(dtype, dask, skipna, func): assert_allclose(actual, expected) -def test_datetime_to_numeric_datetime64(): +@pytest.mark.parametrize("dask", [True, False]) +def test_datetime_to_numeric_datetime64(dask): + if dask and not has_dask: + pytest.skip("requires dask") + times = pd.date_range("2000", periods=5, freq="7D").values - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h") + if dask: + import dask.array + + times = dask.array.from_array(times, chunks=-1) + + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h") expected = 24 * np.arange(0, 35, 7) np.testing.assert_array_equal(result, expected) offset = times[1] - result = duck_array_ops.datetime_to_numeric(times, offset=offset, datetime_unit="h") + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, offset=offset, datetime_unit="h" + ) expected = 24 * np.arange(-7, 28, 7) np.testing.assert_array_equal(result, expected) dtype = np.float32 - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=dtype) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, datetime_unit="h", dtype=dtype + ) expected = 24 * np.arange(0, 35, 7).astype(dtype) np.testing.assert_array_equal(result, expected) @requires_cftime -def test_datetime_to_numeric_cftime(): +@pytest.mark.parametrize("dask", [True, False]) +def test_datetime_to_numeric_cftime(dask): + if dask and not has_dask: + pytest.skip("requires dask") + times = cftime_range("2000", periods=5, freq="7D", calendar="standard").values - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=int) + if dask: + import dask.array + + times = dask.array.from_array(times, chunks=-1) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=int) expected = 24 * np.arange(0, 35, 7) np.testing.assert_array_equal(result, expected) offset = times[1] - result = duck_array_ops.datetime_to_numeric( - times, offset=offset, datetime_unit="h", dtype=int - ) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, offset=offset, datetime_unit="h", dtype=int + ) expected = 24 * np.arange(-7, 28, 7) np.testing.assert_array_equal(result, expected) dtype = np.float32 - result = duck_array_ops.datetime_to_numeric(times, datetime_unit="h", dtype=dtype) + with raise_if_dask_computes(): + result = duck_array_ops.datetime_to_numeric( + times, datetime_unit="h", dtype=dtype + ) expected = 24 * np.arange(0, 35, 7).astype(dtype) np.testing.assert_array_equal(result, expected) diff --git a/xarray/tests/test_formatting.py b/xarray/tests/test_formatting.py index efdb8a57288..a5c044d8ea7 100644 --- a/xarray/tests/test_formatting.py +++ b/xarray/tests/test_formatting.py @@ -9,7 +9,7 @@ import xarray as xr from xarray.core import formatting -from . import requires_netCDF4 +from . import requires_dask, requires_netCDF4 class TestFormatting: @@ -418,6 +418,26 @@ def test_array_repr_variable(self) -> None: with xr.set_options(display_expand_data=False): formatting.array_repr(var) + @requires_dask + def test_array_scalar_format(self) -> None: + var = xr.DataArray(0) + assert var.__format__("") == "0" + assert var.__format__("d") == "0" + assert var.__format__(".2f") == "0.00" + + var = xr.DataArray([0.1, 0.2]) + assert var.__format__("") == "[0.1 0.2]" + with pytest.raises(TypeError) as excinfo: + var.__format__(".2f") + assert "unsupported format string passed to" in str(excinfo.value) + + # also check for dask + var = var.chunk(chunks={"dim_0": 1}) + assert var.__format__("") == "[0.1 0.2]" + with pytest.raises(TypeError) as excinfo: + var.__format__(".2f") + assert "unsupported format string passed to" in str(excinfo.value) + def test_inline_variable_array_repr_custom_repr() -> None: class CustomArray: @@ -480,10 +500,10 @@ def test_short_numpy_repr() -> None: assert num_lines < 30 # threshold option (default: 200) - array = np.arange(100) - assert "..." not in formatting.short_numpy_repr(array) + array2 = np.arange(100) + assert "..." not in formatting.short_numpy_repr(array2) with xr.set_options(display_values_threshold=10): - assert "..." in formatting.short_numpy_repr(array) + assert "..." in formatting.short_numpy_repr(array2) def test_large_array_repr_length() -> None: diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 34f13e1f056..8c745dc640d 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -25,7 +25,10 @@ @pytest.fixture def dataset(): ds = xr.Dataset( - {"foo": (("x", "y", "z"), np.random.randn(3, 4, 2))}, + { + "foo": (("x", "y", "z"), np.random.randn(3, 4, 2)), + "baz": ("x", ["e", "f", "g"]), + }, {"x": ["a", "b", "c"], "y": [1, 2, 3, 4], "z": [1, 2]}, ) ds["boo"] = (("z", "y"), [["f", "g", "h", "j"]] * 2) @@ -72,7 +75,15 @@ def test_multi_index_groupby_map(dataset) -> None: assert_equal(expected, actual) -@pytest.mark.xfail +def test_reduce_numeric_only(dataset) -> None: + gb = dataset.groupby("x", squeeze=False) + with xr.set_options(use_flox=False): + expected = gb.sum() + with xr.set_options(use_flox=True): + actual = gb.sum() + assert_identical(expected, actual) + + def test_multi_index_groupby_sum() -> None: # regression test for GH873 ds = xr.Dataset( @@ -803,47 +814,25 @@ def test_groupby_math_more() -> None: def test_groupby_bins_math(indexed_coord) -> None: N = 7 da = DataArray(np.random.random((N, N)), dims=("x", "y")) - bins = np.arange(0, N + 1, 2) - idxr = DataArray([0, 0, 1, 1, 2, 2], dims="x") - if indexed_coord: da["x"] = np.arange(N) da["y"] = np.arange(N) - idxr["x"] = da["x"][1:] - - g = da.groupby_bins("x", bins) - mean = g.mean() - expected = da.isel(x=slice(1, None)) - mean.isel(x_bins=idxr) - actual = g - mean - assert_identical(expected, actual) - - # non-default cut_kwargs - g = da.groupby_bins("x", bins, right=False) + g = da.groupby_bins("x", np.arange(0, N + 1, 3)) mean = g.mean() - if indexed_coord: - idxr["x"] = da["x"][:-1] - expected = da.isel(x=slice(-1)) - mean.isel(x_bins=idxr) + expected = da.isel(x=slice(1, None)) - mean.isel(x_bins=("x", [0, 0, 0, 1, 1, 1])) actual = g - mean assert_identical(expected, actual) - # mean does not contain all bins in the groupby - actual = g - mean[:-1] - with xr.set_options(arithmetic_join="outer"): - expected = da.isel(x=slice(-1)) - mean.isel(x_bins=idxr) - expected.data[-2:] = np.nan - assert_identical(expected, actual) - def test_groupby_math_nD_group() -> None: N = 40 - chars = ["a", "b", "c", "d", "e", "f", "g", "h"] da = DataArray( np.random.random((N, N)), dims=("x", "y"), coords={ "labels": ( "x", - np.repeat(chars, repeats=N // 8), + np.repeat(["a", "b", "c", "d", "e", "f", "g", "h"], repeats=N // 8), ), }, ) @@ -856,14 +845,6 @@ def test_groupby_math_nD_group() -> None: actual = g - mean assert_identical(expected, actual) - # drop a label from the mean; - # this will require reindexing rather than just a simple sel - actual = g - mean.isel(labels2d=slice(1, -1)) - mean.data[[0, -1]] = np.nan - expected = da - mean.sel(labels2d=da.labels2d) - expected["labels"] = expected.labels.broadcast_like(expected.labels2d) - assert_identical(expected, actual) - da["num"] = ( "x", np.repeat([1, 2, 3, 4, 5, 6, 7, 8], repeats=N // 8), @@ -2008,6 +1989,3 @@ def func(arg1, arg2, arg3=0.0): expected = xr.Dataset({"foo": ("time", [3.0, 3.0, 3.0]), "time": times}) actual = ds.resample(time="D").map(func, args=(1.0,), arg3=1.0) assert_identical(expected, actual) - - -# TODO: move other groupby tests from test_dataset and test_dataarray over here diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index 2a6de0be550..d62254dd327 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -1,10 +1,12 @@ from itertools import combinations, permutations +from typing import cast import numpy as np import pandas as pd import pytest import xarray as xr +from xarray.core.types import InterpOptions from xarray.tests import ( assert_allclose, assert_equal, @@ -24,22 +26,25 @@ pass -def get_example_data(case): - x = np.linspace(0, 1, 100) - y = np.linspace(0, 0.1, 30) - data = xr.DataArray( - np.sin(x[:, np.newaxis]) * np.cos(y), - dims=["x", "y"], - coords={"x": x, "y": y, "x2": ("x", x**2)}, - ) +def get_example_data(case: int) -> xr.DataArray: if case == 0: - return data + # 2D + x = np.linspace(0, 1, 100) + y = np.linspace(0, 0.1, 30) + return xr.DataArray( + np.sin(x[:, np.newaxis]) * np.cos(y), + dims=["x", "y"], + coords={"x": x, "y": y, "x2": ("x", x**2)}, + ) elif case == 1: - return data.chunk({"y": 3}) + # 2D chunked single dim + return get_example_data(0).chunk({"y": 3}) elif case == 2: - return data.chunk({"x": 25, "y": 3}) + # 2D chunked both dims + return get_example_data(0).chunk({"x": 25, "y": 3}) elif case == 3: + # 3D x = np.linspace(0, 1, 100) y = np.linspace(0, 0.1, 30) z = np.linspace(0.1, 0.2, 10) @@ -49,7 +54,10 @@ def get_example_data(case): coords={"x": x, "y": y, "x2": ("x", x**2), "z": z}, ) elif case == 4: + # 3D chunked single dim return get_example_data(3).chunk({"z": 5}) + else: + raise ValueError("case must be 1-4") def test_keywargs(): @@ -62,8 +70,10 @@ def test_keywargs(): @pytest.mark.parametrize("method", ["linear", "cubic"]) @pytest.mark.parametrize("dim", ["x", "y"]) -@pytest.mark.parametrize("case", [0, 1]) -def test_interpolate_1d(method, dim, case): +@pytest.mark.parametrize( + "case", [pytest.param(0, id="no_chunk"), pytest.param(1, id="chunk_y")] +) +def test_interpolate_1d(method: InterpOptions, dim: str, case: int) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -72,7 +82,7 @@ def test_interpolate_1d(method, dim, case): da = get_example_data(case) xdest = np.linspace(0.0, 0.9, 80) - actual = da.interp(method=method, **{dim: xdest}) + actual = da.interp(method=method, coords={dim: xdest}) # scipy interpolation for the reference def func(obj, new_x): @@ -95,7 +105,7 @@ def func(obj, new_x): @pytest.mark.parametrize("method", ["cubic", "zero"]) -def test_interpolate_1d_methods(method): +def test_interpolate_1d_methods(method: InterpOptions) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -103,7 +113,7 @@ def test_interpolate_1d_methods(method): dim = "x" xdest = np.linspace(0.0, 0.9, 80) - actual = da.interp(method=method, **{dim: xdest}) + actual = da.interp(method=method, coords={dim: xdest}) # scipy interpolation for the reference def func(obj, new_x): @@ -122,7 +132,7 @@ def func(obj, new_x): @pytest.mark.parametrize("use_dask", [False, True]) -def test_interpolate_vectorize(use_dask): +def test_interpolate_vectorize(use_dask: bool) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -197,8 +207,10 @@ def func(obj, dim, new_x): assert_allclose(actual, expected.transpose("z", "w", "y", transpose_coords=True)) -@pytest.mark.parametrize("case", [3, 4]) -def test_interpolate_nd(case): +@pytest.mark.parametrize( + "case", [pytest.param(3, id="no_chunk"), pytest.param(4, id="chunked")] +) +def test_interpolate_nd(case: int) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -208,13 +220,13 @@ def test_interpolate_nd(case): da = get_example_data(case) # grid -> grid - xdest = np.linspace(0.1, 1.0, 11) - ydest = np.linspace(0.0, 0.2, 10) - actual = da.interp(x=xdest, y=ydest, method="linear") + xdestnp = np.linspace(0.1, 1.0, 11) + ydestnp = np.linspace(0.0, 0.2, 10) + actual = da.interp(x=xdestnp, y=ydestnp, method="linear") # linear interpolation is separateable - expected = da.interp(x=xdest, method="linear") - expected = expected.interp(y=ydest, method="linear") + expected = da.interp(x=xdestnp, method="linear") + expected = expected.interp(y=ydestnp, method="linear") assert_allclose(actual.transpose("x", "y", "z"), expected.transpose("x", "y", "z")) # grid -> 1d-sample @@ -248,7 +260,7 @@ def test_interpolate_nd(case): @requires_scipy -def test_interpolate_nd_nd(): +def test_interpolate_nd_nd() -> None: """Interpolate nd array with an nd indexer sharing coordinates.""" # Create original array a = [0, 2] @@ -278,7 +290,7 @@ def test_interpolate_nd_nd(): @requires_scipy -def test_interpolate_nd_with_nan(): +def test_interpolate_nd_with_nan() -> None: """Interpolate an array with an nd indexer and `NaN` values.""" # Create indexer into `a` with dimensions (y, x) @@ -298,14 +310,16 @@ def test_interpolate_nd_with_nan(): db = 2 * da ds = xr.Dataset({"da": da, "db": db}) - out = ds.interp(a=ia) + out2 = ds.interp(a=ia) expected_ds = xr.Dataset({"da": expected, "db": 2 * expected}) - xr.testing.assert_allclose(out.drop_vars("a"), expected_ds) + xr.testing.assert_allclose(out2.drop_vars("a"), expected_ds) @pytest.mark.parametrize("method", ["linear"]) -@pytest.mark.parametrize("case", [0, 1]) -def test_interpolate_scalar(method, case): +@pytest.mark.parametrize( + "case", [pytest.param(0, id="no_chunk"), pytest.param(1, id="chunk_y")] +) +def test_interpolate_scalar(method: InterpOptions, case: int) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -333,8 +347,10 @@ def func(obj, new_x): @pytest.mark.parametrize("method", ["linear"]) -@pytest.mark.parametrize("case", [3, 4]) -def test_interpolate_nd_scalar(method, case): +@pytest.mark.parametrize( + "case", [pytest.param(3, id="no_chunk"), pytest.param(4, id="chunked")] +) +def test_interpolate_nd_scalar(method: InterpOptions, case: int) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -361,7 +377,7 @@ def test_interpolate_nd_scalar(method, case): @pytest.mark.parametrize("use_dask", [True, False]) -def test_nans(use_dask): +def test_nans(use_dask: bool) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -377,7 +393,7 @@ def test_nans(use_dask): @pytest.mark.parametrize("use_dask", [True, False]) -def test_errors(use_dask): +def test_errors(use_dask: bool) -> None: if not has_scipy: pytest.skip("scipy is not installed.") @@ -389,7 +405,7 @@ def test_errors(use_dask): for method in ["akima", "spline"]: with pytest.raises(ValueError): - da.interp(x=[0.5, 1.5], method=method) + da.interp(x=[0.5, 1.5], method=method) # type: ignore # not sorted if use_dask: @@ -404,9 +420,9 @@ def test_errors(use_dask): # invalid method with pytest.raises(ValueError): - da.interp(x=[2, 0], method="boo") + da.interp(x=[2, 0], method="boo") # type: ignore with pytest.raises(ValueError): - da.interp(y=[2, 0], method="boo") + da.interp(y=[2, 0], method="boo") # type: ignore # object-type DataArray cannot be interpolated da = xr.DataArray(["a", "b", "c"], dims="x", coords={"x": [0, 1, 2]}) @@ -415,7 +431,7 @@ def test_errors(use_dask): @requires_scipy -def test_dtype(): +def test_dtype() -> None: data_vars = dict( a=("time", np.array([1, 1.25, 2])), b=("time", np.array([True, True, False], dtype=bool)), @@ -432,7 +448,7 @@ def test_dtype(): @requires_scipy -def test_sorted(): +def test_sorted() -> None: # unsorted non-uniform gridded data x = np.random.randn(100) y = np.random.randn(30) @@ -459,7 +475,7 @@ def test_sorted(): @requires_scipy -def test_dimension_wo_coords(): +def test_dimension_wo_coords() -> None: da = xr.DataArray( np.arange(12).reshape(3, 4), dims=["x", "y"], coords={"y": [0, 1, 2, 3]} ) @@ -474,7 +490,7 @@ def test_dimension_wo_coords(): @requires_scipy -def test_dataset(): +def test_dataset() -> None: ds = create_test_data() ds.attrs["foo"] = "var" ds["var1"].attrs["buz"] = "var2" @@ -497,8 +513,8 @@ def test_dataset(): assert interpolated["var1"].attrs["buz"] == "var2" -@pytest.mark.parametrize("case", [0, 3]) -def test_interpolate_dimorder(case): +@pytest.mark.parametrize("case", [pytest.param(0, id="2D"), pytest.param(3, id="3D")]) +def test_interpolate_dimorder(case: int) -> None: """Make sure the resultant dimension order is consistent with .sel()""" if not has_scipy: pytest.skip("scipy is not installed.") @@ -546,7 +562,7 @@ def test_interpolate_dimorder(case): @requires_scipy -def test_interp_like(): +def test_interp_like() -> None: ds = create_test_data() ds.attrs["foo"] = "var" ds["var1"].attrs["buz"] = "var2" @@ -588,7 +604,7 @@ def test_interp_like(): pytest.param("2000-01-01T12:00", 0.5, marks=pytest.mark.xfail), ], ) -def test_datetime(x_new, expected): +def test_datetime(x_new, expected) -> None: da = xr.DataArray( np.arange(24), dims="time", @@ -606,7 +622,7 @@ def test_datetime(x_new, expected): @requires_scipy -def test_datetime_single_string(): +def test_datetime_single_string() -> None: da = xr.DataArray( np.arange(24), dims="time", @@ -620,7 +636,7 @@ def test_datetime_single_string(): @requires_cftime @requires_scipy -def test_cftime(): +def test_cftime() -> None: times = xr.cftime_range("2000", periods=24, freq="D") da = xr.DataArray(np.arange(24), coords=[times], dims="time") @@ -633,7 +649,7 @@ def test_cftime(): @requires_cftime @requires_scipy -def test_cftime_type_error(): +def test_cftime_type_error() -> None: times = xr.cftime_range("2000", periods=24, freq="D") da = xr.DataArray(np.arange(24), coords=[times], dims="time") @@ -646,7 +662,7 @@ def test_cftime_type_error(): @requires_cftime @requires_scipy -def test_cftime_list_of_strings(): +def test_cftime_list_of_strings() -> None: from cftime import DatetimeProlepticGregorian times = xr.cftime_range( @@ -667,7 +683,7 @@ def test_cftime_list_of_strings(): @requires_cftime @requires_scipy -def test_cftime_single_string(): +def test_cftime_single_string() -> None: from cftime import DatetimeProlepticGregorian times = xr.cftime_range( @@ -687,7 +703,7 @@ def test_cftime_single_string(): @requires_scipy -def test_datetime_to_non_datetime_error(): +def test_datetime_to_non_datetime_error() -> None: da = xr.DataArray( np.arange(24), dims="time", @@ -699,7 +715,7 @@ def test_datetime_to_non_datetime_error(): @requires_cftime @requires_scipy -def test_cftime_to_non_cftime_error(): +def test_cftime_to_non_cftime_error() -> None: times = xr.cftime_range("2000", periods=24, freq="D") da = xr.DataArray(np.arange(24), coords=[times], dims="time") @@ -708,7 +724,7 @@ def test_cftime_to_non_cftime_error(): @requires_scipy -def test_datetime_interp_noerror(): +def test_datetime_interp_noerror() -> None: # GH:2667 a = xr.DataArray( np.arange(21).reshape(3, 7), @@ -728,7 +744,7 @@ def test_datetime_interp_noerror(): @requires_cftime @requires_scipy -def test_3641(): +def test_3641() -> None: times = xr.cftime_range("0001", periods=3, freq="500Y") da = xr.DataArray(range(3), dims=["time"], coords=[times]) da.interp(time=["0002-05-01"]) @@ -736,7 +752,7 @@ def test_3641(): @requires_scipy @pytest.mark.parametrize("method", ["nearest", "linear"]) -def test_decompose(method): +def test_decompose(method: InterpOptions) -> None: da = xr.DataArray( np.arange(6).reshape(3, 2), dims=["x", "y"], @@ -769,7 +785,9 @@ def test_decompose(method): for nscalar in range(0, interp_ndim + 1) ], ) -def test_interpolate_chunk_1d(method, data_ndim, interp_ndim, nscalar, chunked): +def test_interpolate_chunk_1d( + method: InterpOptions, data_ndim, interp_ndim, nscalar, chunked: bool +) -> None: """Interpolate nd array with multiple independent indexers It should do a series of 1d interpolation @@ -812,12 +830,15 @@ def test_interpolate_chunk_1d(method, data_ndim, interp_ndim, nscalar, chunked): before = 2 * da.coords[dim][0] - da.coords[dim][1] after = 2 * da.coords[dim][-1] - da.coords[dim][-2] - dest[dim] = np.linspace(before, after, len(da.coords[dim]) * 13) + dest[dim] = cast( + xr.DataArray, + np.linspace(before, after, len(da.coords[dim]) * 13), + ) if chunked: dest[dim] = xr.DataArray(data=dest[dim], dims=[dim]) dest[dim] = dest[dim].chunk(2) - actual = da.interp(method=method, **dest, kwargs=kwargs) - expected = da.compute().interp(method=method, **dest, kwargs=kwargs) + actual = da.interp(method=method, **dest, kwargs=kwargs) # type: ignore + expected = da.compute().interp(method=method, **dest, kwargs=kwargs) # type: ignore assert_identical(actual, expected) @@ -831,7 +852,7 @@ def test_interpolate_chunk_1d(method, data_ndim, interp_ndim, nscalar, chunked): @requires_dask @pytest.mark.parametrize("method", ["linear", "nearest"]) @pytest.mark.filterwarnings("ignore:Increasing number of chunks") -def test_interpolate_chunk_advanced(method): +def test_interpolate_chunk_advanced(method: InterpOptions) -> None: """Interpolate nd array with an nd indexer sharing coordinates.""" # Create original array x = np.linspace(-1, 1, 5) @@ -857,25 +878,25 @@ def test_interpolate_chunk_advanced(method): coords=[("w", w), ("theta", theta)], ) - x = r * np.cos(theta) - y = r * np.sin(theta) - z = xr.DataArray( + xda = r * np.cos(theta) + yda = r * np.sin(theta) + zda = xr.DataArray( data=w[:, np.newaxis] * np.sin(theta), coords=[("w", w), ("theta", theta)], ) kwargs = {"fill_value": None} - expected = da.interp(t=0.5, x=x, y=y, z=z, kwargs=kwargs, method=method) + expected = da.interp(t=0.5, x=xda, y=yda, z=zda, kwargs=kwargs, method=method) da = da.chunk(2) - x = x.chunk(1) - z = z.chunk(3) - actual = da.interp(t=0.5, x=x, y=y, z=z, kwargs=kwargs, method=method) + xda = xda.chunk(1) + zda = zda.chunk(3) + actual = da.interp(t=0.5, x=xda, y=yda, z=zda, kwargs=kwargs, method=method) assert_identical(actual, expected) @requires_scipy -def test_interp1d_bounds_error(): +def test_interp1d_bounds_error() -> None: """Ensure exception on bounds error is raised if requested""" da = xr.DataArray( np.sin(0.3 * np.arange(4)), @@ -898,7 +919,7 @@ def test_interp1d_bounds_error(): (("x", np.array([0, 0.5, 1, 2]), dict(unit="s")), False), ], ) -def test_coord_attrs(x, expect_same_attrs): +def test_coord_attrs(x, expect_same_attrs: bool) -> None: base_attrs = dict(foo="bar") ds = xr.Dataset( data_vars=dict(a=2 * np.arange(5)), @@ -910,7 +931,7 @@ def test_coord_attrs(x, expect_same_attrs): @requires_scipy -def test_interp1d_complex_out_of_bounds(): +def test_interp1d_complex_out_of_bounds() -> None: """Ensure complex nans are used by default""" da = xr.DataArray( np.exp(0.3j * np.arange(4)), diff --git a/xarray/tests/test_sparse.py b/xarray/tests/test_sparse.py index bf4d39105c4..bac1f6407fc 100644 --- a/xarray/tests/test_sparse.py +++ b/xarray/tests/test_sparse.py @@ -7,7 +7,6 @@ from packaging.version import Version import xarray as xr -import xarray.ufuncs as xu from xarray import DataArray, Variable from xarray.core.pycompat import sparse_array_type, sparse_version @@ -279,12 +278,12 @@ def test_unary_op(self): @pytest.mark.filterwarnings("ignore::FutureWarning") def test_univariate_ufunc(self): - assert_sparse_equal(np.sin(self.data), xu.sin(self.var).data) + assert_sparse_equal(np.sin(self.data), np.sin(self.var).data) @pytest.mark.filterwarnings("ignore::FutureWarning") def test_bivariate_ufunc(self): - assert_sparse_equal(np.maximum(self.data, 0), xu.maximum(self.var, 0).data) - assert_sparse_equal(np.maximum(self.data, 0), xu.maximum(0, self.var).data) + assert_sparse_equal(np.maximum(self.data, 0), np.maximum(self.var, 0).data) + assert_sparse_equal(np.maximum(self.data, 0), np.maximum(0, self.var).data) def test_repr(self): expected = dedent( @@ -665,11 +664,6 @@ def test_stack(self): roundtripped = stacked.unstack() assert_identical(arr, roundtripped) - @pytest.mark.filterwarnings("ignore::FutureWarning") - def test_ufuncs(self): - x = self.sp_xr - assert_equal(np.sin(x), xu.sin(x)) - def test_dataarray_repr(self): a = xr.DataArray( sparse.COO.from_numpy(np.ones(4)), diff --git a/xarray/tests/test_testing.py b/xarray/tests/test_testing.py index 2bde7529d1e..1470706d0eb 100644 --- a/xarray/tests/test_testing.py +++ b/xarray/tests/test_testing.py @@ -10,7 +10,7 @@ try: from dask.array import from_array as dask_from_array except ImportError: - dask_from_array = lambda x: x + dask_from_array = lambda x: x # type: ignore try: import pint diff --git a/xarray/tests/test_ufuncs.py b/xarray/tests/test_ufuncs.py index 590ae9ae003..28e5c6cbcb1 100644 --- a/xarray/tests/test_ufuncs.py +++ b/xarray/tests/test_ufuncs.py @@ -1,10 +1,7 @@ -import pickle - import numpy as np import pytest import xarray as xr -import xarray.ufuncs as xu from . import assert_array_equal from . import assert_identical as assert_identical_ @@ -158,52 +155,3 @@ def test_gufuncs(): fake_gufunc = mock.Mock(signature="(n)->()", autospec=np.sin) with pytest.raises(NotImplementedError, match=r"generalized ufuncs"): xarray_obj.__array_ufunc__(fake_gufunc, "__call__", xarray_obj) - - -def test_xarray_ufuncs_deprecation(): - with pytest.warns(FutureWarning, match="xarray.ufuncs"): - xu.cos(xr.DataArray([0, 1])) - - with assert_no_warnings(): - xu.angle(xr.DataArray([0, 1])) - - -@pytest.mark.filterwarnings("ignore::RuntimeWarning") -@pytest.mark.parametrize( - "name", - [ - name - for name in dir(xu) - if ( - not name.startswith("_") - and hasattr(np, name) - and name not in ["print_function", "absolute_import", "division"] - ) - ], -) -def test_numpy_ufuncs(name, request): - x = xr.DataArray([1, 1]) - - np_func = getattr(np, name) - if hasattr(np_func, "nin") and np_func.nin == 2: - args = (x, x) - else: - args = (x,) - - y = np_func(*args) - - if name in ["angle", "iscomplex"]: - # these functions need to be handled with __array_function__ protocol - assert isinstance(y, np.ndarray) - elif name in ["frexp"]: - # np.frexp returns a tuple - assert not isinstance(y, xr.DataArray) - else: - assert isinstance(y, xr.DataArray) - - -@pytest.mark.filterwarnings("ignore:xarray.ufuncs") -def test_xarray_ufuncs_pickle(): - a = 1.0 - cos_pickled = pickle.loads(pickle.dumps(xu.cos)) - assert_identical(cos_pickled(a), xu.cos(a)) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index b8e2f6f4582..886b0360c04 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -2154,6 +2154,40 @@ def test_coarsen_keep_attrs(self, operation="mean"): class TestVariableWithDask(VariableSubclassobjects): cls = staticmethod(lambda *args: Variable(*args).chunk()) + def test_chunk(self): + unblocked = Variable(["dim_0", "dim_1"], np.ones((3, 4))) + assert unblocked.chunks is None + + blocked = unblocked.chunk() + assert blocked.chunks == ((3,), (4,)) + first_dask_name = blocked.data.name + + blocked = unblocked.chunk(chunks=((2, 1), (2, 2))) + assert blocked.chunks == ((2, 1), (2, 2)) + assert blocked.data.name != first_dask_name + + blocked = unblocked.chunk(chunks=(3, 3)) + assert blocked.chunks == ((3,), (3, 1)) + assert blocked.data.name != first_dask_name + + # name doesn't change when rechunking by same amount + # this fails if ReprObject doesn't have __dask_tokenize__ defined + assert unblocked.chunk(2).data.name == unblocked.chunk(2).data.name + + assert blocked.load().chunks is None + + # Check that kwargs are passed + import dask.array as da + + blocked = unblocked.chunk(name="testname_") + assert isinstance(blocked.data, da.Array) + assert "testname_" in blocked.data.name + + # test kwargs form of chunks + blocked = unblocked.chunk(dim_0=3, dim_1=3) + assert blocked.chunks == ((3,), (3, 1)) + assert blocked.data.name != first_dask_name + @pytest.mark.xfail def test_0d_object_array_with_list(self): super().test_0d_object_array_with_list() @@ -2446,7 +2480,7 @@ def test_datetime(self): assert np.ndarray == type(actual) assert np.dtype("datetime64[ns]") == actual.dtype - def test_full_like(self): + def test_full_like(self) -> None: # For more thorough tests, see test_variable.py orig = Variable( dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"} @@ -2469,7 +2503,7 @@ def test_full_like(self): full_like(orig, True, dtype={"x": bool}) @requires_dask - def test_full_like_dask(self): + def test_full_like_dask(self) -> None: orig = Variable( dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"} ).chunk(((1, 1), (2,))) @@ -2500,14 +2534,14 @@ def check(actual, expect_dtype, expect_values): else: assert not isinstance(v, np.ndarray) - def test_zeros_like(self): + def test_zeros_like(self) -> None: orig = Variable( dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"} ) assert_identical(zeros_like(orig), full_like(orig, 0)) assert_identical(zeros_like(orig, dtype=int), full_like(orig, 0, dtype=int)) - def test_ones_like(self): + def test_ones_like(self) -> None: orig = Variable( dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"} ) diff --git a/xarray/ufuncs.py b/xarray/ufuncs.py deleted file mode 100644 index 24907a158ef..00000000000 --- a/xarray/ufuncs.py +++ /dev/null @@ -1,197 +0,0 @@ -"""xarray specific universal functions - -Handles unary and binary operations for the following types, in ascending -priority order: -- scalars -- numpy.ndarray -- dask.array.Array -- xarray.Variable -- xarray.DataArray -- xarray.Dataset -- xarray.core.groupby.GroupBy - -Once NumPy 1.10 comes out with support for overriding ufuncs, this module will -hopefully no longer be necessary. -""" -import textwrap -import warnings as _warnings - -import numpy as _np - -from .core.dataarray import DataArray as _DataArray -from .core.dataset import Dataset as _Dataset -from .core.groupby import GroupBy as _GroupBy -from .core.pycompat import dask_array_type as _dask_array_type -from .core.variable import Variable as _Variable - -_xarray_types = (_Variable, _DataArray, _Dataset, _GroupBy) -_dispatch_order = (_np.ndarray, _dask_array_type) + _xarray_types -_UNDEFINED = object() - - -def _dispatch_priority(obj): - for priority, cls in enumerate(_dispatch_order): - if isinstance(obj, cls): - return priority - return -1 - - -class _UFuncDispatcher: - """Wrapper for dispatching ufuncs.""" - - def __init__(self, name): - self._name = name - - def __call__(self, *args, **kwargs): - if self._name not in ["angle", "iscomplex"]: - _warnings.warn( - "xarray.ufuncs is deprecated. Instead, use numpy ufuncs directly.", - FutureWarning, - stacklevel=2, - ) - - new_args = args - res = _UNDEFINED - if len(args) > 2 or len(args) == 0: - raise TypeError(f"cannot handle {len(args)} arguments for {self._name!r}") - elif len(args) == 1: - if isinstance(args[0], _xarray_types): - res = args[0]._unary_op(self) - else: # len(args) = 2 - p1, p2 = map(_dispatch_priority, args) - if p1 >= p2: - if isinstance(args[0], _xarray_types): - res = args[0]._binary_op(args[1], self) - else: - if isinstance(args[1], _xarray_types): - res = args[1]._binary_op(args[0], self, reflexive=True) - new_args = tuple(reversed(args)) - - if res is _UNDEFINED: - f = getattr(_np, self._name) - res = f(*new_args, **kwargs) - if res is NotImplemented: - raise TypeError( - f"{self._name!r} not implemented for types ({type(args[0])!r}, {type(args[1])!r})" - ) - return res - - -def _skip_signature(doc, name): - if not isinstance(doc, str): - return doc - - if doc.startswith(name): - signature_end = doc.find("\n\n") - doc = doc[signature_end + 2 :] - - return doc - - -def _remove_unused_reference_labels(doc): - if not isinstance(doc, str): - return doc - - max_references = 5 - for num in range(max_references): - label = f".. [{num}]" - reference = f"[{num}]_" - index = f"{num}. " - - if label not in doc or reference in doc: - continue - - doc = doc.replace(label, index) - - return doc - - -def _dedent(doc): - if not isinstance(doc, str): - return doc - - return textwrap.dedent(doc) - - -def _create_op(name): - func = _UFuncDispatcher(name) - func.__name__ = name - doc = getattr(_np, name).__doc__ - - doc = _remove_unused_reference_labels(_skip_signature(_dedent(doc), name)) - - func.__doc__ = ( - f"xarray specific variant of numpy.{name}. Handles " - "xarray.Dataset, xarray.DataArray, xarray.Variable, " - "numpy.ndarray and dask.array.Array objects with " - "automatic dispatching.\n\n" - f"Documentation from numpy:\n\n{doc}" - ) - return func - - -__all__ = ( # noqa: F822 - "angle", - "arccos", - "arccosh", - "arcsin", - "arcsinh", - "arctan", - "arctan2", - "arctanh", - "ceil", - "conj", - "copysign", - "cos", - "cosh", - "deg2rad", - "degrees", - "exp", - "expm1", - "fabs", - "fix", - "floor", - "fmax", - "fmin", - "fmod", - "fmod", - "frexp", - "hypot", - "imag", - "iscomplex", - "isfinite", - "isinf", - "isnan", - "isreal", - "ldexp", - "log", - "log10", - "log1p", - "log2", - "logaddexp", - "logaddexp2", - "logical_and", - "logical_not", - "logical_or", - "logical_xor", - "maximum", - "minimum", - "nextafter", - "rad2deg", - "radians", - "real", - "rint", - "sign", - "signbit", - "sin", - "sinh", - "sqrt", - "square", - "tan", - "tanh", - "trunc", -) - - -for name in __all__: - globals()[name] = _create_op(name) diff --git a/xarray/util/generate_ops.py b/xarray/util/generate_ops.py index f1fd6cbfeb2..0a382642708 100644 --- a/xarray/util/generate_ops.py +++ b/xarray/util/generate_ops.py @@ -210,7 +210,7 @@ def inplace(): try: from dask.array import Array as DaskArray except ImportError: - DaskArray = np.ndarray + DaskArray = np.ndarray # type: ignore # DatasetOpsMixin etc. are parent classes of Dataset etc. # Because of https://github.com/pydata/xarray/issues/5755, we redefine these. Generally diff --git a/xarray/util/generate_reductions.py b/xarray/util/generate_reductions.py index 0b51632343b..96b91c16906 100644 --- a/xarray/util/generate_reductions.py +++ b/xarray/util/generate_reductions.py @@ -33,7 +33,7 @@ try: import flox except ImportError: - flox = None''' + flox = None # type: ignore''' DEFAULT_PREAMBLE = """ @@ -54,6 +54,30 @@ def reduce( GROUPBY_PREAMBLE = """ +class {obj}{cls}Reductions: + _obj: "{obj}" + + def reduce( + self, + func: Callable[..., Any], + dim: Union[None, Hashable, Sequence[Hashable]] = None, + *, + axis: Union[None, int, Sequence[int]] = None, + keep_attrs: bool = None, + keepdims: bool = False, + **kwargs: Any, + ) -> "{obj}": + raise NotImplementedError() + + def _flox_reduce( + self, + dim: Union[None, Hashable, Sequence[Hashable]], + **kwargs, + ) -> "{obj}": + raise NotImplementedError()""" + +RESAMPLE_PREAMBLE = """ + class {obj}{cls}Reductions: _obj: "{obj}" @@ -303,6 +327,7 @@ def generate_code(self, method): extra_kwargs.append(f"numeric_only={method.numeric_only},") # numpy_groupies & flox do not support median + # https://github.com/ml31415/numpy-groupies/issues/43 if method.name == "median": indent = 12 else: @@ -436,7 +461,7 @@ class DataStructure: docref="resampling", docref_description="resampling operations", example_call_preamble='.resample(time="3M")', - definition_preamble=GROUPBY_PREAMBLE, + definition_preamble=RESAMPLE_PREAMBLE, ) DATASET_GROUPBY_GENERATOR = GroupByReductionGenerator( cls="GroupBy", @@ -454,7 +479,7 @@ class DataStructure: docref="resampling", docref_description="resampling operations", example_call_preamble='.resample(time="3M")', - definition_preamble=GROUPBY_PREAMBLE, + definition_preamble=RESAMPLE_PREAMBLE, ) @@ -464,6 +489,7 @@ class DataStructure: p = Path(os.getcwd()) filepath = p.parent / "xarray" / "xarray" / "core" / "_reductions.py" + # filepath = p.parent / "core" / "_reductions.py" # Run from script location with open(filepath, mode="w", encoding="utf-8") as f: f.write(MODULE_PREAMBLE + "\n") for gen in [