rebase with main

Author: Mikejmnez

xarray/backends/pydap_.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import os
-import warnings
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
@@ -37,10 +36,9 @@
 
 
 class PydapArrayWrapper(BackendArray):
-    def __init__(self, array, batch=False, cache=None, checksums=True):
+    def __init__(self, array, batch=None, checksums=True):
         self.array = array
         self._batch = batch
-        self._cache = cache
         self._checksums = checksums
 
     @property
@@ -57,27 +55,19 @@ def __getitem__(self, key):
         )
 
     def _getitem(self, key):
-        if self.array.id in self._cache.keys():
-            # safely avoid re-downloading some coordinates
-            result = self._cache[self.array.id]
-        elif self._batch and hasattr(self.array, "dataset"):
+        if self._batch and hasattr(self.array, "dataset"):
             # this are both True only for pydap>3.5.5
-            from pydap.lib import resolve_batch_for_all_variables
+            from pydap.client import data_check, get_batch_data
 
             dataset = self.array.dataset
-            resolve_batch_for_all_variables(self.array, key, checksums=self._checksums)
-            result = np.asarray(
-                dataset._current_batch_promise.wait_for_result(self.array.id)
-            )
+            get_batch_data(self.array, checksums=self._checksums, key=key)
+            result = data_check(np.asarray(dataset[self.array.id].data), key)
         else:
             result = robust_getitem(self.array, key, catch=ValueError)
-            try:
-                result = np.asarray(result.data)
-            except AttributeError:
-                result = np.asarray(result)
-        axis = tuple(n for n, k in enumerate(key) if isinstance(k, integer_types))
-        if result.ndim + len(axis) != self.array.ndim and axis:
-            result = np.squeeze(result, axis)
+            result = np.asarray(result.data)
+            axis = tuple(n for n, k in enumerate(key) if isinstance(k, integer_types))
+            if result.ndim + len(axis) != self.array.ndim and axis:
+                result = np.squeeze(result, axis)
         return result
 
 
@@ -105,7 +95,7 @@ def __init__(
         dataset,
         group=None,
         session=None,
-        batch=False,
+        batch=None,
         protocol=None,
         checksums=True,
     ):
@@ -119,8 +109,6 @@ def __init__(
         self.dataset = dataset
         self.group = group
         self._batch = batch
-        self._batch_done = False
-        self._array_cache = {}  # holds 1D dimension data
         self._protocol = protocol
         self._checksums = checksums  # true by default
 
@@ -135,7 +123,7 @@ def open(
         timeout=None,
         verify=None,
         user_charset=None,
-        batch=False,
+        batch=None,
         checksums=True,
     ):
         from pydap.client import open_url
@@ -167,34 +155,23 @@ def open(
         elif hasattr(url, "ds"):
             # pydap dataset
             dataset = url.ds
-        args = {"dataset": dataset}
-        args["checksums"] = checksums
+        args = {"dataset": dataset, "checksums": checksums}
         if group:
             args["group"] = group
         if url.startswith(("http", "dap2")):
             args["protocol"] = "dap2"
         elif url.startswith("dap4"):
             args["protocol"] = "dap4"
         if batch:
-            if args["protocol"] == "dap2":
-                warnings.warn(
-                    f"`batch={batch}` is currently only compatible with the `DAP4` "
-                    "protocol. Make sue the OPeNDAP server implements the `DAP4` "
-                    "protocol and then replace the scheme of the url with `dap4` "
-                    "to make use of it. Setting `batch=False`.",
-                    stacklevel=2,
-                )
-            else:
-                # only update if dap4
-                args["batch"] = batch
+            args["batch"] = batch
         return cls(**args)
 
     def open_store_variable(self, var):
-        try:
+        if hasattr(var, "dims"):
             dimensions = [
                 dim.split("/")[-1] if dim.startswith("/") else dim for dim in var.dims
             ]
-        except AttributeError:
+        else:
             # GridType does not have a dims attribute - instead get `dimensions`
             # see https://github.com/pydap/pydap/issues/485
             dimensions = var.dimensions
@@ -214,7 +191,7 @@ def open_store_variable(self, var):
         else:
             # all non-dimension variables
             data = indexing.LazilyIndexedArray(
-                PydapArrayWrapper(var, self._batch, self._array_cache, self._checksums)
+                PydapArrayWrapper(var, self._batch, self._checksums)
             )
 
         return Variable(dimensions, data, var.attributes)
@@ -264,16 +241,14 @@ def _get_data_array(self, var):
         """gets dimension data all at once, storing the numpy
         arrays within a cached dictionary
         """
-        from pydap.lib import get_batch_data
+        from pydap.client import get_batch_data
 
-        if not self._batch_done or var.id not in self._array_cache:
-            # store all dim data into a dict for reuse
-            self._array_cache = get_batch_data(
-                var.parent, self._array_cache, self._checksums
-            )
-            self._batch_done = True
+        if not var._is_data_loaded():
+            # data has not been deserialized yet
+            # runs only once per store/hierarchy
+            get_batch_data(var, checksums=self._checksums)
 
-        return self._array_cache[var.id]
+        return self.dataset[var.id].data
 
 
 class PydapBackendEntrypoint(BackendEntrypoint):
@@ -336,7 +311,7 @@ def open_dataset(
         timeout=None,
         verify=None,
         user_charset=None,
-        batch=False,
+        batch=None,
         checksums=True,
     ) -> Dataset:
         store = PydapDataStore.open(
@@ -382,7 +357,7 @@ def open_datatree(
         timeout=None,
         verify=None,
         user_charset=None,
-        batch=False,
+        batch=None,
         checksums=True,
     ) -> DataTree:
         groups_dict = self.open_groups_as_dict(
@@ -423,7 +398,7 @@ def open_groups_as_dict(
         timeout=None,
         verify=None,
         user_charset=None,
-        batch=False,
+        batch=None,
         checksums=True,
     ) -> dict[str, Dataset]:
         from xarray.core.treenode import NodePath

xarray/tests/test_backends.py

@@ -6564,14 +6564,15 @@ def test_session(self) -> None:
 @network
 @pytest.mark.parametrize("protocol", ["dap2", "dap4"])
 @pytest.mark.parametrize("batch", [False, True])
-def test_batchdap4_downloads(protocol, batch) -> None:
+def test_batchdap4_downloads(tmpdir, protocol, batch) -> None:
     """Test that in dap4, all dimensions are downloaded at once"""
     import pydap
     from pydap.net import create_session
 
     _version_ = Version(pydap.__version__)
     # Create a session with pre-set params in pydap backend, to cache urls
-    session = create_session(use_cache=True, cache_kwargs={"cache_name": "debug"})
+    cache_name = tmpdir / "debug"
+    session = create_session(use_cache=True, cache_kwargs={"cache_name": cache_name})
     session.cache.clear()
     url = "https://test.opendap.org/opendap/hyrax/data/nc/coads_climatology.nc"
 
@@ -6611,25 +6612,6 @@ def test_batchdap4_downloads(protocol, batch) -> None:
         assert len(session.cache.urls()) == 5
 
 
-@requires_pydap
-@network
-def test_batch_warnswithdap2() -> None:
-    from pydap.net import create_session
-
-    # Create a session with pre-set retry params in pydap backend, to cache urls
-    session = create_session(use_cache=True, cache_kwargs={"cache_name": "debug"})
-    session.cache.clear()
-
-    url = "dap2://test.opendap.org/opendap/hyrax/data/nc/coads_climatology.nc"
-    with pytest.warns(UserWarning):
-        open_dataset(
-            url, engine="pydap", session=session, batch=True, decode_times=False
-        )
-
-    # no batching is supported here
-    assert len(session.cache.urls()) == 5
-
-
 class TestEncodingInvalid:
     def test_extract_nc4_variable_encoding(self) -> None:
         var = xr.Variable(("x",), [1, 2, 3], {}, {"foo": "bar"})

xarray/tests/test_backends_datatree.py

@@ -581,7 +581,8 @@ class TestPyDAPDatatreeIO:
     simplegroup_datatree_url = "dap4://test.opendap.org/opendap/dap4/SimpleGroup.nc4.h5"
 
     def test_open_datatree_unaligned_hierarchy(
-        self, url=unaligned_datatree_url
+        self,
+        url=unaligned_datatree_url,
     ) -> None:
         with pytest.raises(
             ValueError,
@@ -614,7 +615,7 @@ def test_open_groups(self, url=unaligned_datatree_url) -> None:
         ) as expected:
             assert_identical(unaligned_dict_of_datasets["/Group1/subgroup1"], expected)
 
-    def test_inherited_coords(self, url=simplegroup_datatree_url) -> None:
+    def test_inherited_coords(self, tmpdir, url=simplegroup_datatree_url) -> None:
         """Test that `open_datatree` inherits coordinates from root tree.
 
         This particular h5 file is a test file that inherits the time coordinate from the root
@@ -644,7 +645,10 @@ def test_inherited_coords(self, url=simplegroup_datatree_url) -> None:
         from pydap.net import create_session
 
         # Create a session with pre-set retry params in pydap backend, to cache urls
-        session = create_session(use_cache=True, cache_kwargs={"cache_name": "debug"})
+        cache_name = tmpdir / "debug"
+        session = create_session(
+            use_cache=True, cache_kwargs={"cache_name": cache_name}
+        )
         session.cache.clear()
 
         _version_ = Version(pydap.__version__)
@@ -661,8 +665,8 @@ def test_inherited_coords(self, url=simplegroup_datatree_url) -> None:
             )
 
         if _version_ > Version("3.5.5"):
-            # Total downloads are: 1 dmr, + 1 dap url for all dimensions across groups
-            assert len(session.cache.urls()) == 2
+            # Total downloads are: 1 dmr, + 1 dap url for all dimensions for each group
+            assert len(session.cache.urls()) == 3
         else:
             # 1 dmr + 1 dap url per dimension (total there are 4 dimension arrays)
             assert len(session.cache.urls()) == 5