code updated after suggestions.

google · Mar 14, 2024 · 41b12e2 · 41b12e2
1 parent d91a9e4
commit 41b12e2
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 17 deletions.
diff --git a/xee/ext.py b/xee/ext.py
@@ -102,6 +102,11 @@ class EarthEngineStore(common.AbstractDataStore):
       'height': 256,
   }
 
+  TILE_FETCH_KWARGS: Dict[str, int] = {
+      'max_retries': 6,
+      'initial_delay': 500,
+  }
+
   SCALE_UNITS: Dict[str, int] = {
       'degree': 1,
       'metre': 10_000,
@@ -147,10 +152,7 @@ def open(
       ee_init_kwargs: Optional[Dict[str, Any]] = None,
       ee_init_if_necessary: bool = False,
       executor_kwargs: Optional[Dict[str, Any]] = None,
-      tile_fetch_kwargs: Dict[str, int] = {
-          'max_retries': 6,
-          'initial_delay': 500,
-      },
+      tile_fetch_kwargs: Dict[str, int] = TILE_FETCH_KWARGS,
   ) -> 'EarthEngineStore':
     if mode != 'r':
       raise ValueError(
@@ -191,10 +193,7 @@ def __init__(
       ee_init_kwargs: Optional[Dict[str, Any]] = None,
       ee_init_if_necessary: bool = False,
       executor_kwargs: Optional[Dict[str, Any]] = None,
-      tile_fetch_kwargs: Dict[str, int] = {
-          'max_retries': 6,
-          'initial_delay': 500,
-      },
+      tile_fetch_kwargs: Dict[str, int] = TILE_FETCH_KWARGS,
   ):
     self.ee_init_kwargs = ee_init_kwargs
     self.ee_init_if_necessary = ee_init_if_necessary
@@ -204,9 +203,7 @@ def __init__(
       executor_kwargs = {}
     self.executor_kwargs = executor_kwargs
 
-    # Default value: (https://github.com/pydata/xarray/blob/main/xarray/backends/common.py#L181).
-    self.tile_fetch_max_retries = tile_fetch_kwargs.pop('max_retries', 6)
-    self.tile_fetch_initial_delay = tile_fetch_kwargs.pop('initial_delay', 500)
+    self.tile_fetch_kwargs = tile_fetch_kwargs
 
     self.image_collection = image_collection
     if n_images != -1:
@@ -498,8 +495,12 @@ def image_to_array(
         pixels_getter,
         params,
         catch=ee.ee_exception.EEException,
-        max_retries=self.tile_fetch_max_retries,
-        initial_delay=self.tile_fetch_initial_delay,
+        max_retries=self.tile_fetch_kwargs.get(
+            'max_retries', self.TILE_FETCH_KWARGS.get('max_retries')
+        ),
+        initial_delay=self.tile_fetch_kwargs.get(
+            'initial_delay', self.TILE_FETCH_KWARGS.get('initial_delay')
+        ),
     )
 
     # Extract out the shape information from EE response.

diff --git a/xee/ext_integration_test.py b/xee/ext_integration_test.py
@@ -81,6 +81,7 @@ def setUp(self):
             '2020-03-30', '2020-04-01'
         ),
         n_images=64,
+        tile_fetch_kwargs={'max_retries': 9},
     )
     self.all_img_store = xee.EarthEngineStore(
         ee.ImageCollection('LANDSAT/LC08/C01/T1').filterDate(
@@ -258,13 +259,16 @@ def __getitem__(self, params):
 
   def test_tile_fetch_kwargs(self):
     arr = xee.EarthEngineBackendArray('B4', self.store)
-    self.assertEqual(arr.store.tile_fetch_initial_delay, 1500)
-    self.assertEqual(arr.store.tile_fetch_max_retries, 10)
+    self.assertEqual(arr.store.tile_fetch_kwargs['initial_delay'], 1500)
+    self.assertEqual(arr.store.tile_fetch_kwargs['max_retries'], 10)
 
     arr1 = xee.EarthEngineBackendArray('longitude', self.lnglat_store)
-    self.assertEqual(arr1.store.tile_fetch_initial_delay, 500)
-    self.assertEqual(arr1.store.tile_fetch_max_retries, 6)
+    self.assertEqual(arr1.store.tile_fetch_kwargs['initial_delay'], 500)
+    self.assertEqual(arr1.store.tile_fetch_kwargs['max_retries'], 6)
 
+    arr1 = xee.EarthEngineBackendArray('spi2y', self.conus_store)
+    self.assertNotIn('initial_delay', arr1.store.tile_fetch_kwargs)
+    self.assertEqual(arr1.store.tile_fetch_kwargs['max_retries'], 9)
 
 class EEBackendEntrypointTest(absltest.TestCase):