fix netcdf tests

Author: b8raoult

climetlab/readers/grib/index/__init__.py

@@ -44,7 +44,7 @@ def __init__(self, *args, **kwargs):
 
     @classmethod
     def new_mask_index(self, *args, **kwargs):
-        return MaskFieldSet(*args, **kwargs)
+        return GribMaskFieldSet(*args, **kwargs)
 
     @property
     def availability_path(self):
@@ -53,7 +53,7 @@ def availability_path(self):
     @classmethod
     def merge(cls, sources):
         assert all(isinstance(_, GribFieldSet) for _ in sources)
-        return MultiFieldSet(sources)
+        return GribMultiFieldSet(sources)
 
     def available(self, request, as_list_of_dicts=False):
         from climetlab.utils.availability import Availability
@@ -152,12 +152,12 @@ def _normalize_kwargs_names(self, **kwargs):
         return kwargs
 
 
-class MaskFieldSet(GribFieldSet, MaskIndex):
+class GribMaskFieldSet(GribFieldSet, MaskIndex):
     def __init__(self, *args, **kwargs):
         MaskIndex.__init__(self, *args, **kwargs)
 
 
-class MultiFieldSet(GribFieldSet, MultiIndex):
+class GribMultiFieldSet(GribFieldSet, MultiIndex):
     def __init__(self, *args, **kwargs):
         MultiIndex.__init__(self, *args, **kwargs)
 

climetlab/readers/grib/output.py

@@ -199,9 +199,9 @@ def update_metadata(self, handle, metadata, compulsary):
         if "number" in metadata:
             compulsary += ("numberOfForecastsInEnsemble",)
             productDefinitionTemplateNumber = {"tp": 11}
-            metadata["productDefinitionTemplateNumber"] = (
-                productDefinitionTemplateNumber.get(handle.get("shortName"), 1)
-            )
+            metadata[
+                "productDefinitionTemplateNumber"
+            ] = productDefinitionTemplateNumber.get(handle.get("shortName"), 1)
 
         if metadata.get("type") in ("pf", "cf"):
             metadata.setdefault("typeOfGeneratingProcess", 4)

climetlab/readers/grib/reader.py

@@ -10,7 +10,7 @@
 import logging
 
 from climetlab.readers import Reader
-from climetlab.readers.grib.index import MultiFieldSet
+from climetlab.readers.grib.index import GribMultiFieldSet
 from climetlab.readers.grib.index.file import FieldSetInOneFile
 
 LOG = logging.getLogger(__name__)
@@ -31,7 +31,7 @@ def merge(cls, readers):
         assert all(isinstance(s, GRIBReader) for s in readers), readers
         assert len(readers) > 1
 
-        return MultiFieldSet(readers)
+        return GribMultiFieldSet(readers)
 
     def mutate_source(self):
         # A GRIBReader is a source itself

climetlab/readers/netcdf/fieldset.py

@@ -10,6 +10,7 @@
 from functools import cached_property
 from itertools import product
 
+from climetlab.core.index import MaskIndex, MultiIndex
 from climetlab.indexing.fieldset import FieldSet
 from climetlab.utils.bbox import BoundingBox
 from climetlab.utils.dates import to_datetime
@@ -24,6 +25,10 @@ def __init__(self, path):
         self.path = path
         self.opendap = path.startswith("http")
 
+    @classmethod
+    def new_mask_index(self, *args, **kwargs):
+        return NetCDFMaskFieldSet(*args, **kwargs)
+
     def __repr__(self):
         return "NetCDFReader(%s)" % (self.path,)
 
@@ -40,6 +45,9 @@ def __getitem__(self, n):
     def dataset(self):
         import xarray as xr
 
+        if ".zarr" in self.path:
+            return xr.open_zarr(self.path)
+
         if self.opendap:
             return xr.open_dataset(self.path)
         else:
@@ -146,7 +154,7 @@ def _get_fields(self, ds):  # noqa C901
     def to_xarray(self, **kwargs):
         import xarray as xr
 
-        if self.opendap:
+        if self.path.startswith("http"):
             return xr.open_dataset(self.path, **kwargs)
         return type(self).to_xarray_multi_from_paths([self.path], **kwargs)
 
@@ -185,3 +193,31 @@ def to_datetime_list(self):
 
     def to_bounding_box(self):
         return BoundingBox.multi_merge([s.to_bounding_box() for s in self.fields])
+
+    @classmethod
+    def merge(cls, sources):
+        assert len(sources) > 1
+        assert all(isinstance(_, NetCDFFieldSet) for _ in sources)
+        return NetCDFMultiFieldSet(sources)
+
+
+class NetCDFMaskFieldSet(NetCDFFieldSet, MaskIndex):
+    def __init__(self, *args, **kwargs):
+        MaskIndex.__init__(self, *args, **kwargs)
+
+
+class NetCDFMultiFieldSet(NetCDFFieldSet, MultiIndex):
+    def __init__(self, *args, **kwargs):
+        MultiIndex.__init__(self, *args, **kwargs)
+        self.paths = [s.path for s in args[0]]
+
+    def to_xarray(self, **kwargs):
+        import xarray as xr
+        if not kwargs:
+            kwargs = dict(combine="by_coords")
+        return xr.open_mfdataset(self.paths, **kwargs)
+
+
+    @cached_property
+    def dataset(self):
+        return self.to_xarray(combine="by_coords")

climetlab/sources/indexed_urls.py

@@ -10,7 +10,7 @@
 import warnings
 
 from climetlab.indexing import PerUrlIndex
-from climetlab.readers.grib.index import MultiFieldSet
+from climetlab.readers.grib.index import GribMultiFieldSet
 from climetlab.readers.grib.index.sql import FieldsetInFilesWithSqlIndex
 from climetlab.sources.indexed import IndexedSource
 from climetlab.utils.patterns import Pattern
@@ -60,7 +60,7 @@ def __init__(
             # This is to avoid keeping them on the request
             request.pop(used)
 
-        index = MultiFieldSet(
+        index = GribMultiFieldSet(
             FieldsetInFilesWithSqlIndex.from_url(
                 get_index_url(url, substitute_extension, index_extension),
                 selection=request,

climetlab/utils/patterns.py

@@ -88,7 +88,14 @@ def substitute(self, value, name):
         return self.format % value
 
 
-TYPES = {"": Any, "int": Int, "float": Float, "date": Datetime, "strftime": Datetime, "enum": Enum}
+TYPES = {
+    "": Any,
+    "int": Int,
+    "float": Float,
+    "date": Datetime,
+    "strftime": Datetime,
+    "enum": Enum,
+}
 
 
 class Constant:

tests/readers/test_netcdf_reader.py

@@ -97,7 +97,6 @@ def test_dummy_netcdf_4():
 @pytest.mark.long_test
 @pytest.mark.download
 @pytest.mark.skipif(NO_CDS, reason="No access to CDS")
-@pytest.mark.skipif(True, reason="Merging of netcdf files does not work yet")
 def test_multi():
     s1 = load_source(
         "cds",

tests/sources/test_merge.py

@@ -72,7 +72,7 @@ def test_nc_merge_custom(custom_merger):
     target2 = xr.open_mfdataset([s1.path, s2.path])
     assert target2.identical(merged)
 
-@pytest.mark.skipif(True, reason="Merging of netcdf files does not work yet")
+
 def test_nc_merge_var():
     s1 = load_source(
         "climetlab-testing",
@@ -125,7 +125,7 @@ def _merge_var_different_coords(kind1, kind2):
 
     assert target.identical(merged)
 
-@pytest.mark.skipif(True, reason="Merging of netcdf files does not work yet")
+
 def test_nc_merge_var_different_coords():
     _merge_var_different_coords("netcdf", "netcdf")