add xarray-datatree and test

ci/environment-py310.yml

@@ -7,6 +7,7 @@ dependencies:
   - dask
   - zarr
   - xarray
+  - xarray-datatree
   - h5netcdf
   - h5py<3.9
   - pandas

ci/environment-py38.yml

@@ -7,6 +7,7 @@ dependencies:
   - dask
   - zarr
   - xarray
+  - xarray-datatree
   - h5netcdf
   - h5py<3.9
   - pandas

ci/environment-py39.yml

@@ -7,6 +7,7 @@ dependencies:
   - dask
   - zarr
   - xarray
+  - xarray-datatree
   - h5netcdf
   - h5py<3.9
   - pandas

kerchunk/grib2.py

@@ -23,6 +23,7 @@
 
 from kerchunk.utils import class_factory, _encode_for_JSON
 from kerchunk.codecs import GRIBCodec
+from kerchunk.combine import MultiZarrToZarr, drop
 
 
 # cfgrib copies over certain GRIB attributes
@@ -334,8 +335,6 @@ def example_combine(
     ...        "consolidated": False,
     ...        "storage_options": {"fo": tot, "remote_options": {"anon": True}}})
     """
-    from kerchunk.combine import MultiZarrToZarr, drop
-
     files = [
         "s3://noaa-hrrr-bdp-pds/hrrr.20190101/conus/hrrr.t22z.wrfsfcf01.grib2",
         "s3://noaa-hrrr-bdp-pds/hrrr.20190101/conus/hrrr.t23z.wrfsfcf01.grib2",
@@ -391,8 +390,6 @@ def grib_tree(
     list(dict): A new zarr store like dictionary for use as a reference filesystem mapper with zarr
     or xarray datatree
     """
-    from kerchunk.combine import MultiZarrToZarr
-
     # Hard code the filters in the correct order for the group hierarchy
     filters = ["stepType", "typeOfLevel"]
 
@@ -526,8 +523,8 @@ def grib_tree(
 
 def correct_hrrr_subhf_step(group: Dict) -> Dict:
     """
-    Overrides the definition of the "step" variable. 
-    
+    Overrides the definition of the "step" variable.
+
     Sets the value equal to the `valid_time - time`
     in hours as a floating point value. This fixes issues with the HRRR SubHF grib2 step as read by
     cfgrib via scan_grib.

kerchunk/tests/test_grib.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pytest
 import xarray as xr
+import datatree
 import zarr
 import ujson
 from kerchunk.grib2 import (
@@ -192,7 +193,7 @@ def test_hrrr_subhf_corrected_grib_tree():
     assert zg.dswrf.avg.surface.time.shape == (1,)
 
 
-# The following test use json fixture data generated from calling scan grib
+# The following two test use json fixture data generated from calling scan grib
 #   scan_grib("testdata/hrrr.t01z.wrfsfcf00.grib2")
 #   scan_grib("testdata/hrrr.t01z.wrfsfcf01.grib2")
 # and filtering the results for keys starting with "dswrf" or "u"
@@ -248,3 +249,21 @@ def test_hrrr_sfcf_grib_tree():
 
     assert zg.u.instant.isobaricInhPa.time[:].tolist() == [1695862800]
     assert zg.u.instant.isobaricInhPa.time.shape == (1,)
+
+
+def test_hrrr_sfcf_grib_datatree():
+    fpath = os.path.join(here, "hrrr.wrfsfcf.subset.json")
+    with open(fpath, "rb") as fobj:
+        scanned_msgs = ujson.load(fobj)
+    merged = grib_tree(scanned_msgs)
+    dt = datatree.open_datatree(
+        fsspec.filesystem("reference", fo=merged).get_mapper(""),
+        engine="zarr",
+        consolidated=False,
+    )
+    # Assert a few things... but if it loads we are mostly done.
+    np.testing.assert_array_equal(
+        dt.u.instant.heightAboveGround.step.values[:],
+        np.array([0, 3600 * 10**9], dtype="timedelta64[ns]"),
+    )
+    assert dt.u.attrs == dict(name="U component of wind")

pyproject.toml

@@ -44,6 +44,7 @@ dev = [
     "s3fs",
     "types-ujson",
     "xarray",
+    "xarray-datatree",
     "cfgrib",
     "scipy",
     "netcdf4"