fix netcdf tests

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")