Automatically support NCZarr

Author: shoyer

xarray/backends/api.py

@@ -99,16 +99,16 @@
 
 
 def get_default_netcdf_write_engine(
+    path_or_file: str | IOBase | None,
     format: T_NetcdfTypes | None,
-    to_fileobject: bool,
 ) -> Literal["netcdf4", "h5netcdf", "scipy"]:
     """Return the default netCDF library to use for writing a netCDF file."""
+
     module_names = {
         "netcdf4": "netCDF4",
         "scipy": "scipy",
         "h5netcdf": "h5netcdf",
     }
-
     candidates = list(OPTIONS["netcdf_engine_order"])
 
     if format is not None:
@@ -128,15 +128,24 @@ def get_default_netcdf_write_engine(
         if format not in {"NETCDF3_64BIT", "NETCDF3_CLASSIC"}:
             candidates.remove("scipy")
 
-    if to_fileobject:
+    nczarr_mode = isinstance(path_or_file, str) and path_or_file.endswith(
+        "#mode=nczarr"
+    )
+    if nczarr_mode:
+        candidates[:] = ["netcdf4"]
+
+    if isinstance(path_or_file, IOBase):
         candidates.remove("netcdf4")
 
     for engine in candidates:
         module_name = module_names[engine]
         if importlib.util.find_spec(module_name) is not None:
             return engine
 
-    format_str = f" with {format=}" if format is not None else ""
+    if nczarr_mode:
+        format_str = " in NCZarr format"
+    else:
+        format_str = f" with {format=}" if format is not None else ""
     libraries = ", ".join(module_names[c] for c in candidates)
     raise ValueError(
         f"cannot write NetCDF files{format_str} because none of the suitable "
@@ -2077,8 +2086,7 @@ def to_netcdf(
     path_or_file = _normalize_path(path_or_file)
 
     if engine is None:
-        to_fileobject = isinstance(path_or_file, IOBase)
-        engine = get_default_netcdf_write_engine(format, to_fileobject)
+        engine = get_default_netcdf_write_engine(path_or_file, format)
 
     # validate Dataset keys, DataArray names, and attr keys/values
     _validate_dataset_names(dataset)

xarray/core/datatree_io.py

@@ -55,10 +55,9 @@ def _datatree_to_netcdf(
     filepath = _normalize_path(filepath)
 
     if engine is None:
-        to_fileobject = isinstance(filepath, io.IOBase)
         engine = get_default_netcdf_write_engine(
+            path_or_file=filepath,
             format="NETCDF4",  # required for supporting groups
-            to_fileobject=to_fileobject,
         )  # type: ignore[assignment]
 
     if group is not None:

xarray/tests/test_backends.py

@@ -7277,7 +7277,8 @@ def _create_nczarr(self, filename):
         # https://github.com/Unidata/netcdf-c/issues/2259
         ds = ds.drop_vars("dim3")
 
-        ds.to_netcdf(f"file://{filename}#mode=nczarr", engine="netcdf4")
+        # engine="netcdf4" is not required for backwards compatibility
+        ds.to_netcdf(f"file://{filename}#mode=nczarr")
         return ds
 
     def test_open_nczarr(self) -> None:

xarray/tests/test_backends_api.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import io
 import re
 import sys
 from numbers import Number
@@ -23,53 +24,44 @@
 @requires_scipy
 @requires_h5netcdf
 def test_get_default_netcdf_write_engine() -> None:
-    engine = get_default_netcdf_write_engine(format=None, to_fileobject=False)
+    engine = get_default_netcdf_write_engine("", format=None)
     assert engine == "h5netcdf"
 
-    engine = get_default_netcdf_write_engine(format="NETCDF4", to_fileobject=False)
+    engine = get_default_netcdf_write_engine("", format="NETCDF4")
     assert engine == "h5netcdf"
 
-    engine = get_default_netcdf_write_engine(
-        format="NETCDF4_CLASSIC", to_fileobject=False
-    )
+    engine = get_default_netcdf_write_engine("", format="NETCDF4_CLASSIC")
     assert engine == "netcdf4"
 
-    engine = get_default_netcdf_write_engine(format="NETCDF4", to_fileobject=True)
+    engine = get_default_netcdf_write_engine(io.BytesIO(), format="NETCDF4")
     assert engine == "h5netcdf"
 
-    engine = get_default_netcdf_write_engine(
-        format="NETCDF3_CLASSIC", to_fileobject=False
-    )
+    engine = get_default_netcdf_write_engine("", format="NETCDF3_CLASSIC")
     assert engine == "scipy"
 
-    engine = get_default_netcdf_write_engine(
-        format="NETCDF3_CLASSIC", to_fileobject=True
-    )
+    engine = get_default_netcdf_write_engine(io.BytesIO(), format="NETCDF3_CLASSIC")
     assert engine == "scipy"
 
+    engine = get_default_netcdf_write_engine("path.zarr#mode=nczarr", format=None)
+    assert engine == "netcdf4"
+
     with xr.set_options(netcdf_engine_order=["netcdf4", "scipy", "h5netcdf"]):
-        engine = get_default_netcdf_write_engine(format=None, to_fileobject=False)
+        engine = get_default_netcdf_write_engine("", format=None)
         assert engine == "netcdf4"
 
-        engine = get_default_netcdf_write_engine(format="NETCDF4", to_fileobject=False)
+        engine = get_default_netcdf_write_engine("", format="NETCDF4")
         assert engine == "netcdf4"
 
-        engine = get_default_netcdf_write_engine(
-            format="NETCDF4_CLASSIC", to_fileobject=False
-        )
+        engine = get_default_netcdf_write_engine("", format="NETCDF4_CLASSIC")
         assert engine == "netcdf4"
 
-        engine = get_default_netcdf_write_engine(format="NETCDF4", to_fileobject=True)
+        engine = get_default_netcdf_write_engine(io.BytesIO(), format="NETCDF4")
         assert engine == "h5netcdf"
 
-        engine = get_default_netcdf_write_engine(
-            format="NETCDF3_CLASSIC", to_fileobject=False
-        )
+        engine = get_default_netcdf_write_engine("", format="NETCDF3_CLASSIC")
         assert engine == "netcdf4"
 
-        engine = get_default_netcdf_write_engine(
-            format="NETCDF3_CLASSIC", to_fileobject=True
-        )
+        engine = get_default_netcdf_write_engine(io.BytesIO(), format="NETCDF3_CLASSIC")
         assert engine == "scipy"
 
 
@@ -81,17 +73,31 @@ def test_default_engine_h5netcdf(monkeypatch):
     monkeypatch.delitem(sys.modules, "scipy", raising=False)
     monkeypatch.setattr(sys, "meta_path", [])
 
-    engine = get_default_netcdf_write_engine(format=None, to_fileobject=False)
+    engine = get_default_netcdf_write_engine("", format=None)
     assert engine == "h5netcdf"
 
     with pytest.raises(
         ValueError,
         match=re.escape(
             "cannot write NetCDF files with format='NETCDF3_CLASSIC' because "
-            "none of the suitable backend libraries (netCDF4, scipy) are installed"
+            "none of the suitable backend libraries (scipy, netCDF4) are installed"
+        ),
+    ):
+        get_default_netcdf_write_engine("", format="NETCDF3_CLASSIC")
+
+
+def test_default_engine_nczarr_no_netcdf4_python(monkeypatch):
+    monkeypatch.delitem(sys.modules, "netCDF4", raising=False)
+    monkeypatch.setattr(sys, "meta_path", [])
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "cannot write NetCDF files in NCZarr format because "
+            "none of the suitable backend libraries (netCDF4) are installed"
         ),
     ):
-        get_default_netcdf_write_engine(format="NETCDF3_CLASSIC", to_fileobject=False)
+        get_default_netcdf_write_engine("#mode=nczarr", format=None)
 
 
 def test_custom_engine() -> None: