Merge branch 'main' of https://github.com/pydata/xarray

TomNicholas · TomNicholas · commit 0b8fa41fdd2e · 2025-03-20T10:28:01.000-04:00
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # 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)
+[![CI](https://github.com/pydata/xarray/actions/workflows/ci.yaml/badge.svg?branch=main)](https://github.com/pydata/xarray/actions/workflows/ci.yaml?query=branch%3Amain)
 [![Code coverage](https://codecov.io/gh/pydata/xarray/branch/main/graph/badge.svg?flag=unittests)](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://asv-runner.github.io/asv-collection/xarray/)
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -14,11 +14,47 @@ What's New
 
     np.random.seed(123456)
 
-.. _whats-new.2025.02.0:
+.. _whats-new.2025.03.1:
 
-v2025.02.0 (unreleased)
+v2025.03.1 (unreleased)
 -----------------------
 
+New Features
+~~~~~~~~~~~~
+
+
+Breaking changes
+~~~~~~~~~~~~~~~~
+
+
+Deprecations
+~~~~~~~~~~~~
+
+
+Bug fixes
+~~~~~~~~~
+
+
+Documentation
+~~~~~~~~~~~~~
+
+
+Internal Changes
+~~~~~~~~~~~~~~~~
+
+.. _whats-new.2025.03.0:
+
+v2025.03.0 (Mar 20, 2025)
+-------------------------
+
+This release brings tested support for Python 3.13, support for reading Zarr V3 datasets into a :py:class:`~xarray.DataTree`,
+significant improvements to datetime & timedelta encoding/decoding, and improvements to the :py:class:`~xarray.DataTree` API;
+in addition to the usual bug fixes and other improvements.
+Thanks to the 26 contributors to this release:
+Alfonso Ladino, Benoit Bovy, Chuck Daniels, Deepak Cherian, Eni, Florian Jetter, Ian Hunt-Isaak, Jan, Joe Hamman, Josh Kihm, Julia Signell,
+Justus Magin, Kai Mühlbauer, Kobe Vandelanotte, Mathias Hauser, Max Jones, Maximilian Roos, Oliver Watt-Meyer, Sam Levang, Sander van Rijn,
+Spencer Clark, Stephan Hoyer, Tom Nicholas, Tom White, Vecko and maddogghoek
+
 New Features
 ~~~~~~~~~~~~
 - Added :py:meth:`tutorial.open_datatree` and :py:meth:`tutorial.load_datatree`
@@ -66,6 +102,10 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+
+- Fix ``open_datatree`` incompatibilities with Zarr-Python V3 and refactor
+  ``TestZarrDatatreeIO`` accordingly (:issue:`9960`, :pull:`10020`).
+  By `Alfonso Ladino-Rincon <https://github.com/aladinor>`_.
 - Default to resolution-dependent optimal integer encoding units when saving
   chunked non-nanosecond :py:class:`numpy.datetime64` or
   :py:class:`numpy.timedelta64` arrays to disk. Previously units of
@@ -75,8 +115,8 @@ Bug fixes
 - Use mean of min/max years as offset in calculation of datetime64 mean
   (:issue:`10019`, :pull:`10035`).
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
-- Fix DataArray().drop_attrs(deep=False) and add support for attrs to
-  DataArray()._replace(). (:issue:`10027`, :pull:`10030`). By `Jan
+- Fix ``DataArray().drop_attrs(deep=False)`` and add support for attrs to
+  ``DataArray()._replace()``. (:issue:`10027`, :pull:`10030`). By `Jan
   Haacker <https://github.com/j-haacker>`_.
 - Fix bug preventing encoding times with missing values with small integer
   dtype (:issue:`9134`, :pull:`9498`). By `Spencer Clark
@@ -97,14 +137,12 @@ Bug fixes
   datetimes and timedeltas (:issue:`8957`, :pull:`10050`).
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
 
+
 Documentation
 ~~~~~~~~~~~~~
 - Better expose the :py:class:`Coordinates` class in API reference (:pull:`10000`)
   By `Benoit Bovy <https://github.com/benbovy>`_.
 
-Internal Changes
-~~~~~~~~~~~~~~~~
-
 
 .. _whats-new.2025.01.2:
 
@@ -179,9 +217,6 @@ New Features
   :py:class:`pandas.DatetimeIndex` (:pull:`9965`). By `Spencer Clark
   <https://github.com/spencerkclark>`_ and `Kai Mühlbauer
   <https://github.com/kmuehlbauer>`_.
-
-Breaking changes
-~~~~~~~~~~~~~~~~
 - Adds shards to the list of valid_encodings in the zarr backend, so that
   sharded Zarr V3s can be written (:issue:`9947`, :pull:`9948`).
   By `Jacob Prince_Bieker <https://github.com/jacobbieker>`_
diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py
@@ -666,10 +666,21 @@ def open_store(
             use_zarr_fill_value_as_mask=use_zarr_fill_value_as_mask,
             zarr_format=zarr_format,
         )
+
+        from zarr import Group
+
+        group_members: dict[str, Group] = {}
         group_paths = list(_iter_zarr_groups(zarr_group, parent=group))
-        return {
+        for path in group_paths:
+            if path == group:
+                group_members[path] = zarr_group
+            else:
+                rel_path = path.removeprefix(f"{group}/")
+                group_members[path] = zarr_group[rel_path.removeprefix("/")]
+
+        out = {
             group: cls(
-                zarr_group.get(group),
+                group_store,
                 mode,
                 consolidate_on_close,
                 append_dim,
@@ -680,8 +691,9 @@ def open_store(
                 use_zarr_fill_value_as_mask,
                 cache_members=cache_members,
             )
-            for group in group_paths
+            for group, group_store in group_members.items()
         }
+        return out
 
     @classmethod
     def open_group(
@@ -1034,8 +1046,6 @@ def store(
         if self._consolidate_on_close:
             kwargs = {}
             if _zarr_v3():
-                # https://github.com/zarr-developers/zarr-python/pull/2113#issuecomment-2386718323
-                kwargs["path"] = self.zarr_group.name.lstrip("/")
                 kwargs["zarr_format"] = self.zarr_group.metadata.zarr_format
             zarr.consolidate_metadata(self.zarr_group.store, **kwargs)
 
@@ -1662,8 +1672,6 @@ def open_groups_as_dict(
         zarr_version=None,
         zarr_format=None,
     ) -> dict[str, Dataset]:
-        from xarray.core.treenode import NodePath
-
         filename_or_obj = _normalize_path(filename_or_obj)
 
         # Check for a group and make it a parent if it exists
@@ -1686,7 +1694,6 @@ def open_groups_as_dict(
         )
 
         groups_dict = {}
-
         for path_group, store in stores.items():
             store_entrypoint = StoreBackendEntrypoint()
 
@@ -1762,44 +1769,57 @@ def _get_open_params(
             consolidated = False
 
     if _zarr_v3():
-        missing_exc = ValueError
+        # TODO: replace AssertionError after https://github.com/zarr-developers/zarr-python/issues/2821 is resolved
+        missing_exc = AssertionError
     else:
         missing_exc = zarr.errors.GroupNotFoundError
 
-    if consolidated is None:
-        try:
-            zarr_group = zarr.open_consolidated(store, **open_kwargs)
-        except (ValueError, KeyError):
-            # ValueError in zarr-python 3.x, KeyError in 2.x.
+    if consolidated in [None, True]:
+        # open the root of the store, in case there is metadata consolidated there
+        group = open_kwargs.pop("path")
+
+        if consolidated:
+            # TODO: an option to pass the metadata_key keyword
+            zarr_root_group = zarr.open_consolidated(store, **open_kwargs)
+        elif consolidated is None:
+            # same but with more error handling in case no consolidated metadata found
             try:
-                zarr_group = zarr.open_group(store, **open_kwargs)
-                emit_user_level_warning(
-                    "Failed to open Zarr store with consolidated metadata, "
-                    "but successfully read with non-consolidated metadata. "
-                    "This is typically much slower for opening a dataset. "
-                    "To silence this warning, consider:\n"
-                    "1. Consolidating metadata in this existing store with "
-                    "zarr.consolidate_metadata().\n"
-                    "2. Explicitly setting consolidated=False, to avoid trying "
-                    "to read consolidate metadata, or\n"
-                    "3. Explicitly setting consolidated=True, to raise an "
-                    "error in this case instead of falling back to try "
-                    "reading non-consolidated metadata.",
-                    RuntimeWarning,
-                )
-            except missing_exc as err:
-                raise FileNotFoundError(
-                    f"No such file or directory: '{store}'"
-                ) from err
-    elif consolidated:
-        # TODO: an option to pass the metadata_key keyword
-        zarr_group = zarr.open_consolidated(store, **open_kwargs)
+                zarr_root_group = zarr.open_consolidated(store, **open_kwargs)
+            except (ValueError, KeyError):
+                # ValueError in zarr-python 3.x, KeyError in 2.x.
+                try:
+                    zarr_root_group = zarr.open_group(store, **open_kwargs)
+                    emit_user_level_warning(
+                        "Failed to open Zarr store with consolidated metadata, "
+                        "but successfully read with non-consolidated metadata. "
+                        "This is typically much slower for opening a dataset. "
+                        "To silence this warning, consider:\n"
+                        "1. Consolidating metadata in this existing store with "
+                        "zarr.consolidate_metadata().\n"
+                        "2. Explicitly setting consolidated=False, to avoid trying "
+                        "to read consolidate metadata, or\n"
+                        "3. Explicitly setting consolidated=True, to raise an "
+                        "error in this case instead of falling back to try "
+                        "reading non-consolidated metadata.",
+                        RuntimeWarning,
+                    )
+                except missing_exc as err:
+                    raise FileNotFoundError(
+                        f"No such file or directory: '{store}'"
+                    ) from err
+
+        # but the user should still receive a DataTree whose root is the group they asked for
+        if group and group != "/":
+            zarr_group = zarr_root_group[group.removeprefix("/")]
+        else:
+            zarr_group = zarr_root_group
     else:
         if _zarr_v3():
             # we have determined that we don't want to use consolidated metadata
             # so we set that to False to avoid trying to read it
             open_kwargs["use_consolidated"] = False
         zarr_group = zarr.open_group(store, **open_kwargs)
+
     close_store_on_close = zarr_group.store is not store
 
     # we use this to determine how to handle fill_value
diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py
@@ -16,7 +16,6 @@
     TYPE_CHECKING,
     Any,
     Concatenate,
-    Literal,
     NoReturn,
     ParamSpec,
     TypeVar,
@@ -1741,7 +1740,7 @@ def to_zarr(
         consolidated: bool = True,
         group: str | None = None,
         write_inherited_coords: bool = False,
-        compute: Literal[True] = True,
+        compute: bool = True,
         **kwargs,
     ):
         """
diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py
@@ -16,6 +16,7 @@
 
 import xarray.testing
 from xarray import Dataset
+from xarray.coding.times import _STANDARD_CALENDARS as _STANDARD_CALENDARS_UNSORTED
 from xarray.core.duck_array_ops import allclose_or_equiv  # noqa: F401
 from xarray.core.extension_array import PandasExtensionArray
 from xarray.core.options import set_options
@@ -164,6 +165,21 @@ def _importorskip(
 
 has_array_api_strict, requires_array_api_strict = _importorskip("array_api_strict")
 
+parametrize_zarr_format = pytest.mark.parametrize(
+    "zarr_format",
+    [
+        pytest.param(2, id="zarr_format=2"),
+        pytest.param(
+            3,
+            marks=pytest.mark.skipif(
+                not has_zarr_v3,
+                reason="zarr-python v2 cannot understand the zarr v3 format",
+            ),
+            id="zarr_format=3",
+        ),
+    ],
+)
+
 
 def _importorskip_h5netcdf_ros3(has_h5netcdf: bool):
     if not has_h5netcdf:
@@ -355,13 +371,36 @@ def create_test_data(
     return obj
 
 
-_CFTIME_CALENDARS = [
+_STANDARD_CALENDAR_NAMES = sorted(_STANDARD_CALENDARS_UNSORTED)
+_NON_STANDARD_CALENDAR_NAMES = {
+    "noleap",
     "365_day",
     "360_day",
     "julian",
     "all_leap",
     "366_day",
-    "gregorian",
-    "proleptic_gregorian",
-    "standard",
+}
+_NON_STANDARD_CALENDARS = [
+    pytest.param(cal, marks=requires_cftime)
+    for cal in sorted(_NON_STANDARD_CALENDAR_NAMES)
 ]
+_STANDARD_CALENDARS = [pytest.param(cal) for cal in _STANDARD_CALENDAR_NAMES]
+_ALL_CALENDARS = sorted(_STANDARD_CALENDARS + _NON_STANDARD_CALENDARS)
+_CFTIME_CALENDARS = [
+    pytest.param(*p.values, marks=requires_cftime) for p in _ALL_CALENDARS
+]
+
+
+def _all_cftime_date_types():
+    import cftime
+
+    return {
+        "noleap": cftime.DatetimeNoLeap,
+        "365_day": cftime.DatetimeNoLeap,
+        "360_day": cftime.Datetime360Day,
+        "julian": cftime.DatetimeJulian,
+        "all_leap": cftime.DatetimeAllLeap,
+        "366_day": cftime.DatetimeAllLeap,
+        "gregorian": cftime.DatetimeGregorian,
+        "proleptic_gregorian": cftime.DatetimeProlepticGregorian,
+    }
diff --git a/xarray/tests/test_accessor_dt.py b/xarray/tests/test_accessor_dt.py
@@ -6,6 +6,8 @@
 
 import xarray as xr
 from xarray.tests import (
+    _CFTIME_CALENDARS,
+    _all_cftime_date_types,
     assert_allclose,
     assert_array_equal,
     assert_chunks_equal,
@@ -390,15 +392,6 @@ def test_dask_accessor_method(self, method, parameters) -> None:
         assert_equal(actual.compute(), expected.compute())
 
 
-_CFTIME_CALENDARS = [
-    "365_day",
-    "360_day",
-    "julian",
-    "all_leap",
-    "366_day",
-    "gregorian",
-    "proleptic_gregorian",
-]
 _NT = 100
 
 
@@ -407,6 +400,13 @@ def calendar(request):
     return request.param
 
 
+@pytest.fixture()
+def cftime_date_type(calendar):
+    if calendar == "standard":
+        calendar = "proleptic_gregorian"
+    return _all_cftime_date_types()[calendar]
+
+
 @pytest.fixture()
 def times(calendar):
     import cftime
@@ -573,13 +573,6 @@ def test_dask_field_access(times_3d, data, field) -> None:
     assert_equal(result.compute(), expected)
 
 
-@pytest.fixture()
-def cftime_date_type(calendar):
-    from xarray.tests.test_coding_times import _all_cftime_date_types
-
-    return _all_cftime_date_types()[calendar]
-
-
 @requires_cftime
 def test_seasons(cftime_date_type) -> None:
     dates = xr.DataArray(
diff --git a/xarray/tests/test_backends_datatree.py b/xarray/tests/test_backends_datatree.py
diff --git a/xarray/tests/test_cftimeindex.py b/xarray/tests/test_cftimeindex.py
diff --git a/xarray/tests/test_coding_times.py b/xarray/tests/test_coding_times.py
