Repository: ecmwf/earthkit-data
Pull request #165 (2023-08-29 14:54:35)

Implement xarray fieldlist
Fixes to #158.

Implements xarray fieldlists. Example:

```
>>> import earthkit.data
>>> import xarray as xr
>>> ds = earthkit.data.from_object(xr.open_dataset("docs/examples/tuv_pl.nc"))
>>> ds
<earthkit.data.readers.netcdf.XArrayFieldList object at 0x16059c490>
>>> ds.ls()
   variable  level                time    units
0         t   1000 2018-08-01 12:00:00        K
1         t    850 2018-08-01 12:00:00        K
2         t    700 2018-08-01 12:00:00        K
3         t    500 2018-08-01 12:00:00        K
4         t    400 2018-08-01 12:00:00        K
5         t    300 2018-08-01 12:00:00        K
6         u   1000 2018-08-01 12:00:00  m s**-1
7         u    850 2018-08-01 12:00:00  m s**-1
8         u    700 2018-08-01 12:00:00  m s**-1
9         u    500 2018-08-01 12:00:00  m s**-1
10        u    400 2018-08-01 12:00:00  m s**-1
11        u    300 2018-08-01 12:00:00  m s**-1
12        v   1000 2018-08-01 12:00:00  m s**-1
13        v    850 2018-08-01 12:00:00  m s**-1
14        v    700 2018-08-01 12:00:00  m s**-1
15        v    500 2018-08-01 12:00:00  m s**-1
16        v    400 2018-08-01 12:00:00  m s**-1
17        v    300 2018-08-01 12:00:00  m s**-1
>>> ds[0]
XArrayField(t,level=1000,time=2018-08-01 12:00:00)
>>> ds[0].metadata("level")
1000
```


Patch:
diff --git a/earthkit/data/readers/netcdf.py b/earthkit/data/readers/netcdf.py
index a4bac3f..20c7cac 100644
--- a/earthkit/data/readers/netcdf.py
+++ b/earthkit/data/readers/netcdf.py
@@ -144,7 +144,7 @@ class DataSet:
         return self._bbox[(lat, lon)]
 
 
-class NetCDFFieldGeography(Geography):
+class XArrayFieldGeography(Geography):
     def __init__(self, metadata, da, ds, variable):
         self.metadata = metadata
         self._da = da
@@ -209,11 +209,11 @@ class NetCDFFieldGeography(Geography):
         return grid_mapping
 
 
-class NetCDFMetadata(RawMetadata):
+class XArrayMetadata(RawMetadata):
     def __init__(self, field):
-        if not isinstance(field, NetCDFField):
+        if not isinstance(field, XArrayField):
             raise TypeError(
-                f"NetCDFMetadata: expected field type NetCDFField, got {type(field)}"
+                f"XArrayMetadata: expected field type XArrayField, got {type(field)}"
             )
         self._field = field
         self._geo = None
@@ -233,7 +233,7 @@ class NetCDFMetadata(RawMetadata):
     @property
     def geography(self):
         if self._geo is None:
-            self._geo = NetCDFFieldGeography(
+            self._geo = XArrayFieldGeography(
                 self, self._field._da, self._field._ds, self._field.variable
             )
         return self._geo
@@ -247,7 +247,7 @@ class NetCDFMetadata(RawMetadata):
         return LS_KEYS
 
 
-class NetCDFField(Field):
+class XArrayField(Field):
     def __init__(self, ds, variable, slices, non_dim_coords):
         super().__init__()
 
@@ -288,13 +288,13 @@ class NetCDFField(Field):
 
     def __repr__(self):
         return (
-            f"NetCDFField({self.variable},"
+            f"{self.__class__.__name__}({self.variable},"
             + ",".join([f"{s.name}={s.value}" for s in self.slices])
             + ")"
         )
 
     def _make_metadata(self):
-        return NetCDFMetadata(self)
+        return XArrayMetadata(self)
 
     def to_xarray(self):
         dims = self._da.dims
@@ -327,9 +327,11 @@ class NetCDFField(Field):
         return values
 
 
-class NetCDFFieldList(FieldList):
-    def __init__(self, path, *args, **kwargs):
-        self.path = path
+class XArrayFieldListCore(FieldList):
+    FIELD_TYPE = None
+
+    def __init__(self, ds, *args, **kwargs):
+        self.ds = ds
         self._fields = None
         Index.__init__(self, *args, **kwargs)
 
@@ -344,12 +346,7 @@ class NetCDFFieldList(FieldList):
             self._fields = self._get_fields()
 
     def _get_fields(self):
-        import xarray as xr
-
-        with closing(
-            xr.open_mfdataset(self.path, combine="by_coords")
-        ) as ds:  # or nested
-            return self._get_fields_from_ds(DataSet(ds))
+        return self._get_fields_from_ds(DataSet(self.ds))
 
     def _get_fields_from_ds(self, ds):  # noqa C901
         # Select only geographical variables
@@ -450,13 +447,95 @@ class NetCDFFieldList(FieldList):
                 for value, coordinate in zip(values, coordinates):
                     slices.append(coordinate.make_slice(value))
 
-                fields.append(NetCDFField(ds, name, slices, non_dim_coords))
+                fields.append(self.FIELD_TYPE(ds, name, slices, non_dim_coords))
 
         # if not fields:
         #     raise Exception("NetCDFReader no 2D fields found in %s" % (self.path,))
 
         return fields
 
+    def to_pandas(self):
+        return self.to_xarray().to_pandas()
+
+    def to_xarray(self, **kwargs):
+        return self.ds
+
+    def to_netcdf(self, *args, **kwargs):
+        """
+        Save the data to a netCDF file.
+
+        Parameters
+        ----------
+        See `xarray.DataArray.to_netcdf`.
+        """
+        return self.ds.to_netcdf(*args, **kwargs)
+
+    @classmethod
+    def merge(cls, sources):
+        assert all(isinstance(_, XArrayFieldList) for _ in sources)
+        return XArrayMultiFieldList(sources)
+
+    @classmethod
+    def new_mask_index(cls, *args, **kwargs):
+        return XArrayMaskFieldList(*args, **kwargs)
+
+
+class XArrayFieldList(XArrayFieldListCore):
+    VERSION = 1
+
+    def __init__(self, ds, **kwargs):
+        self.FIELD_TYPE = XArrayField
+        super().__init__(ds, **kwargs)
+
+    def __getitem__(self, n):
+        if isinstance(n, int):
+            return self.fields[n]
+        else:
+            return super().__getitem__(n)
+
+    def __len__(self):
+        return len(self.fields)
+
+
+class XArrayMaskFieldList(XArrayFieldListCore, MaskIndex):
+    def __init__(self, *args, **kwargs):
+        MaskIndex.__init__(self, *args, **kwargs)
+
+
+class XArrayMultiFieldList(XArrayFieldListCore, MultiIndex):
+    def __init__(self, *args, **kwargs):
+        MultiIndex.__init__(self, *args, **kwargs)
+
+    def to_xarray(self, **kwargs):
+        import xarray as xr
+
+        return xr.merge([x.ds for x in self.indexes], **kwargs)
+
+
+class NetCDFMetadata(XArrayMetadata):
+    pass
+
+
+class NetCDFField(XArrayField):
+    def _make_metadata(self):
+        return NetCDFMetadata(self)
+
+
+class NetCDFFieldList(XArrayFieldListCore):
+    def __init__(self, path, *args, **kwargs):
+        self.FIELD_TYPE = NetCDFField
+        self.path = path
+        # self._fields = None
+        super().__init__(None, *args, **kwargs)
+
+    def _get_fields(self):
+        import xarray as xr
+
+        with closing(
+            xr.open_mfdataset(self.path, combine="by_coords")
+        ) as ds:  # or nested
+            return self._get_fields_from_ds(DataSet(ds))
+
     @classmethod
     def merge(cls, sources):
         assert all(isinstance(_, NetCDFFieldList) for _ in sources)
@@ -466,12 +545,19 @@ class NetCDFFieldList(FieldList):
     def new_mask_index(self, *args, **kwargs):
         return NetCDFMaskFieldList(*args, **kwargs)
 
-    def to_pandas(self):
-        return self.to_xarray().to_pandas()
-
     def to_xarray(self, **kwargs):
         return type(self).to_xarray_multi_from_paths(self.path, **kwargs)
 
+    def to_netcdf(self, *args, **kwargs):
+        """
+        Save the data to a netCDF file.
+
+        Parameters
+        ----------
+        See `xarray.DataArray.to_netcdf`.
+        """
+        return self.to_xarray().to_netcdf(*args, **kwargs)
+
     @classmethod
     def to_xarray_multi_from_paths(cls, paths, **kwargs):
         import xarray as xr
@@ -532,38 +618,10 @@ class NetCDFFieldListReader(NetCDFFieldListInOneFile, Reader):
     def __repr__(self):
         return "NetCDFFieldListReader(%s)" % (self.path,)
 
-    # @classmethod
-    # def merge(cls, readers):
-    #     assert all(isinstance(s, NetCDFFieldListReader) for s in readers), readers
-    #     assert len(readers) > 1
-
-    #     return NetCDFFieldListReader(readers[0], [s.path for s in readers])
-
     def mutate_source(self):
         # A NetCDFReader is a source itself
         return self
 
-    # def to_pandas(self):
-    #     return self.to_xarray().to_pandas()
-
-    # def to_xarray(self, **kwargs):
-    #     return type(self).to_xarray_multi_from_paths(self.path, **kwargs)
-
-    # @classmethod
-    # def to_xarray_multi_from_paths(cls, paths, **kwargs):
-    #     import xarray as xr
-
-    #     if not isinstance(paths, list):
-    #         paths = [paths]
-
-    #     options = dict()
-    #     options.update(kwargs.get("xarray_open_mfdataset_kwargs", {}))
-
-    #     return xr.open_mfdataset(
-    #         paths,
-    #         **options,
-    #     )
-
 
 class NetCDFReader(Reader):
     def __init__(self, source, path):
@@ -610,8 +668,8 @@ def _match_magic(magic, deeper_check):
 
 def reader(source, path, magic=None, deeper_check=False):
     if _match_magic(magic, deeper_check):
-        r = NetCDFFieldListReader(source, path)
-        if len(r) > 0:
-            return r
+        fs = NetCDFFieldListReader(source, path)
+        if len(fs) > 0:
+            return fs
         else:
             return NetCDFReader(source, path)
diff --git a/earthkit/data/wrappers/xarray.py b/earthkit/data/wrappers/xarray.py
index 0cfa202..4d088ea 100644
--- a/earthkit/data/wrappers/xarray.py
+++ b/earthkit/data/wrappers/xarray.py
@@ -9,6 +9,7 @@
 
 # from emohawk.metadata import AXES, COMPONENTS
 
+from earthkit.data.readers import netcdf
 from earthkit.data.wrappers import Wrapper
 
 
@@ -133,13 +134,20 @@ class XArrayDatasetWrapper(XArrayDataArrayWrapper):
 def wrapper(data, *args, **kwargs):
     import xarray as xr
 
+    ds = None
     if isinstance(data, xr.Dataset):
-        return XArrayDatasetWrapper(data, *args, **kwargs)
-
-    if isinstance(data, xr.DataArray):
+        ds = data
+    elif isinstance(data, xr.DataArray):
         try:
-            return XArrayDatasetWrapper(data.to_dataset(), *args, **kwargs)
+            ds = data.to_dataset()
         except ValueError:
             return XArrayDataArrayWrapper(data, *args, **kwargs)
 
+    if ds is not None:
+        fs = netcdf.XArrayFieldList(ds, **kwargs)
+        if len(fs) > 0:
+            return fs
+        else:
+            return XArrayDatasetWrapper(ds, *args, **kwargs)
+
     return None


Test patch:
diff --git a/earthkit/data/testing.py b/earthkit/data/testing.py
index 4f2055e..d192be3 100644
--- a/earthkit/data/testing.py
+++ b/earthkit/data/testing.py
@@ -14,7 +14,7 @@ from contextlib import contextmanager
 from importlib import import_module
 from unittest.mock import patch
 
-from earthkit.data import from_source
+from earthkit.data import from_object, from_source
 from earthkit.data.readers.text import TextReader
 from earthkit.data.sources.empty import EmptySource
 
@@ -126,6 +126,15 @@ def check_unsafe_archives(extension):
         check(ds)
 
 
+def load_nc_or_xr_source(path, mode):
+    if mode == "nc":
+        return from_source("file", path)
+    else:
+        import xarray
+
+        return from_object(xarray.open_dataset(path))
+
+
 def main(path):
     import sys
 
diff --git a/tests/netcdf/test_netcdf_concat.py b/tests/netcdf/test_netcdf_concat.py
index c2f8cab..0b9776f 100644
--- a/tests/netcdf/test_netcdf_concat.py
+++ b/tests/netcdf/test_netcdf_concat.py
@@ -11,13 +11,15 @@
 
 import datetime
 
-from earthkit.data import from_source
-from earthkit.data.testing import earthkit_test_data_file
+import pytest
 
+from earthkit.data.testing import earthkit_test_data_file, load_nc_or_xr_source
 
-def test_netcdf_concat():
-    ds1 = from_source("file", earthkit_test_data_file("era5_2t_1.nc"))
-    ds2 = from_source("file", earthkit_test_data_file("era5_2t_2.nc"))
+
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_concat(mode):
+    ds1 = load_nc_or_xr_source(earthkit_test_data_file("era5_2t_1.nc"), mode)
+    ds2 = load_nc_or_xr_source(earthkit_test_data_file("era5_2t_2.nc"), mode)
     ds = ds1 + ds2
 
     assert len(ds) == 2
diff --git a/tests/netcdf/test_netcdf_metadata.py b/tests/netcdf/test_netcdf_metadata.py
index f604c7c..60c206d 100644
--- a/tests/netcdf/test_netcdf_metadata.py
+++ b/tests/netcdf/test_netcdf_metadata.py
@@ -14,9 +14,10 @@ import datetime
 import pytest
 
 from earthkit.data import from_source
-from earthkit.data.testing import earthkit_examples_file
+from earthkit.data.testing import earthkit_examples_file, load_nc_or_xr_source
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize(
     "key,expected_value",
     [
@@ -34,8 +35,9 @@ from earthkit.data.testing import earthkit_examples_file
         (("variable", "level"), ("t", 1000)),
     ],
 )
-def test_netcdf_metadata_single_field(key, expected_value):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_metadata_single_field(mode, key, expected_value):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
+
     # sn = f.metadata(key)
     # assert sn == [expected_value]
     sn = f[0].metadata(key)
@@ -77,8 +79,9 @@ def test_netcdf_datetime():
     assert ds.datetime() == ref
 
 
-def test_netcdf_valid_datetime():
-    ds = from_source("file", earthkit_examples_file("test.nc"))
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_valid_datetime(mode):
+    ds = load_nc_or_xr_source(earthkit_examples_file("test.nc"), mode)
     assert ds[0].metadata("valid_datetime") == datetime.datetime(2020, 5, 13, 12)
 
 
diff --git a/tests/netcdf/test_netcdf_sel.py b/tests/netcdf/test_netcdf_sel.py
index f4112a9..e41c34e 100644
--- a/tests/netcdf/test_netcdf_sel.py
+++ b/tests/netcdf/test_netcdf_sel.py
@@ -13,10 +13,10 @@ import datetime
 
 import pytest
 
-from earthkit.data import from_source
-from earthkit.data.testing import earthkit_examples_file
+from earthkit.data.testing import earthkit_examples_file, load_nc_or_xr_source
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize(
     "params,expected_meta,metadata_keys",
     [
@@ -47,8 +47,8 @@ from earthkit.data.testing import earthkit_examples_file
         ),
     ],
 )
-def test_netcdf_sel_single_file_1(params, expected_meta, metadata_keys):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_sel_single_file_1(mode, params, expected_meta, metadata_keys):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
 
     g = f.sel(**params)
     assert len(g) == len(expected_meta)
diff --git a/tests/netcdf/test_netcdf_slice.py b/tests/netcdf/test_netcdf_slice.py
index e5d6f05..7ee5c86 100644
--- a/tests/netcdf/test_netcdf_slice.py
+++ b/tests/netcdf/test_netcdf_slice.py
@@ -13,9 +13,10 @@ import numpy as np
 import pytest
 
 from earthkit.data import from_source
-from earthkit.data.testing import earthkit_examples_file
+from earthkit.data.testing import earthkit_examples_file, load_nc_or_xr_source
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize(
     "index,expected_meta",
     [
@@ -26,8 +27,8 @@ from earthkit.data.testing import earthkit_examples_file
         (-5, ["v", 850]),
     ],
 )
-def test_netcdf_single_index(index, expected_meta):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_single_index(mode, index, expected_meta):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
 
     r = f[index]
     assert r.metadata(["variable", "level"]) == expected_meta
@@ -43,6 +44,7 @@ def test_netcdf_single_index_bad():
         f[27]
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize(
     "indexes,expected_meta",
     [
@@ -54,8 +56,8 @@ def test_netcdf_single_index_bad():
         (slice(14, None), [["v", 700], ["v", 500], ["v", 400], ["v", 300]]),
     ],
 )
-def test_netcdf_slice_single_file(indexes, expected_meta):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_slice_single_file(mode, indexes, expected_meta):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
     r = f[indexes]
     assert len(r) == 4
     assert r.metadata(["variable", "level"]) == expected_meta
@@ -91,12 +93,14 @@ def test_netcdf_slice_multi_file(indexes, expected_meta):
     assert f.metadata("shortName") == ["2t", "msl", "t", "z", "t", "z"]
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize(
     "indexes1,indexes2",
     [(np.array([1, 16, 5, 9]), np.array([1, 3])), ([1, 16, 5, 9], [1, 3])],
 )
-def test_netcdf_array_indexing(indexes1, indexes2):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_array_indexing(mode, indexes1, indexes2):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
+
     r = f[indexes1]
     assert len(r) == 4
     assert r.metadata("variable") == ["t", "v", "t", "u"]
@@ -106,15 +110,17 @@ def test_netcdf_array_indexing(indexes1, indexes2):
     assert r1.metadata("variable") == ["v", "u"]
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize("indexes", [(np.array([1, 19, 5, 9])), ([1, 19, 5, 9])])
-def test_netcdf_array_indexing_bad(indexes):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_array_indexing_bad(mode, indexes):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
     with pytest.raises(IndexError):
         f[indexes]
 
 
-def test_netcdf_fieldlist_iterator():
-    g = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_fieldlist_iterator(mode):
+    g = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
     sn = g.metadata("variable")
     assert len(sn) == 18
     iter_sn = [f.metadata("variable") for f in g]
@@ -155,8 +161,9 @@ def test_netcdf_fieldlist_iterator_with_zip_multiple():
         assert levs2 == ref_levs, i
 
 
-def test_netcdf_fieldlist_reverse_iterator():
-    g = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_fieldlist_reverse_iterator(mode):
+    g = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
     sn = g.metadata("variable")
     sn_reversed = list(reversed(sn))
     assert sn_reversed[0] == "v"
diff --git a/tests/netcdf/test_netcdf_summary.py b/tests/netcdf/test_netcdf_summary.py
index c988fa8..aef884b 100644
--- a/tests/netcdf/test_netcdf_summary.py
+++ b/tests/netcdf/test_netcdf_summary.py
@@ -10,12 +10,14 @@
 
 import datetime
 
-from earthkit.data import from_source
-from earthkit.data.testing import earthkit_examples_file
+import pytest
 
+from earthkit.data.testing import earthkit_examples_file, load_nc_or_xr_source
 
-def test_netcdf_ls():
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_ls(mode):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
 
     # default keys
     f1 = f[:4]
diff --git a/tests/netcdf/test_netcdf_values.py b/tests/netcdf/test_netcdf_values.py
index a33772b..7944b91 100644
--- a/tests/netcdf/test_netcdf_values.py
+++ b/tests/netcdf/test_netcdf_values.py
@@ -12,8 +12,7 @@
 import numpy as np
 import pytest
 
-from earthkit.data import from_source
-from earthkit.data.testing import earthkit_examples_file
+from earthkit.data.testing import earthkit_examples_file, load_nc_or_xr_source
 
 
 def check_array(v, shape=None, first=None, last=None, meanv=None, eps=1e-3):
@@ -23,8 +22,9 @@ def check_array(v, shape=None, first=None, last=None, meanv=None, eps=1e-3):
     assert np.isclose(v.mean(), meanv, eps)
 
 
-def test_netcdf_values_surf():
-    f = from_source("file", earthkit_examples_file("test.nc"))
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_values_surf(mode):
+    f = load_nc_or_xr_source(earthkit_examples_file("test.nc"), mode)
 
     eps = 1e-5
 
@@ -59,8 +59,9 @@ def test_netcdf_values_surf():
     assert np.allclose(v0_f, v0, eps)
 
 
-def test_netcdf_values_upper():
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_values_upper(mode):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
 
     eps = 1e-5
 
@@ -89,8 +90,9 @@ def test_netcdf_values_upper():
     )
 
 
-def test_netcdf_to_numpy_surf():
-    f = from_source("file", earthkit_examples_file("test.nc"))
+@pytest.mark.parametrize("mode", ["nc", "xr"])
+def test_netcdf_to_numpy_surf(mode):
+    f = load_nc_or_xr_source(earthkit_examples_file("test.nc"), mode)
 
     eps = 1e-5
     v = f.to_numpy()
@@ -124,6 +126,7 @@ def test_netcdf_to_numpy_surf():
     assert np.allclose(v0_f.flatten(), v0, eps)
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize(
     "first,options, expected_shape",
     [
@@ -135,8 +138,8 @@ def test_netcdf_to_numpy_surf():
         (True, {"flatten": False}, (11, 19)),
     ],
 )
-def test_netcdf_to_numpy_surf_shape(first, options, expected_shape):
-    f = from_source("file", earthkit_examples_file("test.nc"))
+def test_netcdf_to_numpy_surf_shape(mode, first, options, expected_shape):
+    f = load_nc_or_xr_source(earthkit_examples_file("test.nc"), mode)
 
     eps = 1e-5
 
@@ -154,6 +157,7 @@ def test_netcdf_to_numpy_surf_shape(first, options, expected_shape):
     assert np.allclose(v_ref, v1, eps)
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize(
     "options, expected_shape",
     [
@@ -175,8 +179,8 @@ def test_netcdf_to_numpy_surf_shape(first, options, expected_shape):
         ({"flatten": False}, (18, 7, 12)),
     ],
 )
-def test_netcdf_to_numpy_upper_shape(options, expected_shape):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_to_numpy_upper_shape(mode, options, expected_shape):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
 
     eps = 1e-5
 
@@ -198,9 +202,10 @@ def test_netcdf_to_numpy_upper_shape(options, expected_shape):
     assert np.allclose(vf15, vr, eps)
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-def test_netcdf_to_numpy_surf_dtype(dtype):
-    f = from_source("file", earthkit_examples_file("test.nc"))
+def test_netcdf_to_numpy_surf_dtype(mode, dtype):
+    f = load_nc_or_xr_source(earthkit_examples_file("test.nc"), mode)
 
     v = f[0].to_numpy(dtype=dtype)
     assert v.dtype == dtype
@@ -209,9 +214,10 @@ def test_netcdf_to_numpy_surf_dtype(dtype):
     assert v.dtype == dtype
 
 
+@pytest.mark.parametrize("mode", ["nc", "xr"])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-def test_netcdf_to_numpy_upper_dtype(dtype):
-    f = from_source("file", earthkit_examples_file("tuv_pl.nc"))
+def test_netcdf_to_numpy_upper_dtype(mode, dtype):
+    f = load_nc_or_xr_source(earthkit_examples_file("tuv_pl.nc"), mode)
 
     v = f[0].to_numpy(dtype=dtype)
     assert v.dtype == dtype