add n_tries parameter

Author: malmans2

tests/test_20_open_dataset.py

@@ -4,7 +4,7 @@
 import xarray as xr
 
 
-@pytest.mark.parametrize("download", [False])
+@pytest.mark.parametrize("download", [True, False])
 def test_open_dataset(tmp_path: Path, index_node: str, download: bool) -> None:
     esgpull_path = tmp_path / "esgpull"
     selection = {

xarray_esgf/client.py

@@ -44,6 +44,7 @@ class Client:
     selection: dict[str, str | list[str]]
     esgpull_path: str | Path | None = None
     index_node: str | None = None
+    n_tries: int = 1
 
     @cached_property
     def _client(self) -> Esgpull:
@@ -75,9 +76,15 @@ def missing_files(self) -> list[File]:
         ]
 
     def download(self) -> list[File]:
-        downloaded, errors = asyncio.run(
-            self._client.download(self.missing_files, use_db=False)
-        )
+        files = []
+        for _ in range(self.n_tries):
+            downloaded, errors = asyncio.run(
+                self._client.download(self.missing_files, use_db=False)
+            )
+            files.extend(downloaded)
+            if not errors:
+                break
+
         exceptions = []
         for error in errors:
             err = error.err
@@ -86,7 +93,7 @@ def download(self) -> list[File]:
         if exceptions:
             msg = "Download errors"
             raise ExceptionGroup(msg, exceptions)
-        return downloaded
+        return files
 
     @use_new_combine_kwarg_defaults
     def open_dataset(

xarray_esgf/engine.py

@@ -19,11 +19,13 @@ def open_dataset(  # type: ignore[override]
         concat_dims: DATASET_ID_KEYS | Iterable[DATASET_ID_KEYS] | None = None,
         download: bool = False,
         show_progress: bool = True,
+        n_tries: int = 1,
     ) -> Dataset:
         client = Client(
             selection=filename_or_obj,
             esgpull_path=esgpull_path,
             index_node=index_node,
+            n_tries=n_tries,
         )
         return client.open_dataset(
             concat_dims=concat_dims,