refactor use of earthaccess

Author: kzalite

HydroEO/satellites/swot/__init__.py

@@ -2,8 +2,6 @@
 
 import datetime
 
-import earthaccess
-
 from HydroEO.satellites.swot import download as _download
 from HydroEO.satellites.swot import preprocess as _preprocess
 
@@ -21,15 +19,13 @@ def query(
         startdate=startdate,
         enddate=enddate,
         product=product,
-        earthaccess_client=earthaccess,
     )
 
 
 def download(results, download_directory: str):
     return _download.download(
         results=results,
         download_directory=download_directory,
-        earthaccess_client=earthaccess,
     )
 
 

HydroEO/satellites/swot/download.py

@@ -1,3 +1,4 @@
+import contextlib
 import datetime
 import logging
 import os
@@ -15,18 +16,123 @@
 SWOT_LAKE_SHORT_NAME = "SWOT_L2_HR_LakeSP_D"
 
 
+@contextlib.contextmanager
+def _suppress_granule_size_warning():
+    """Suppress known earthaccess DataGranule.size deprecation warning."""
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r"As of version 1\.0, `DataGranule\.size` will be accessed as an attribute",
+            category=FutureWarning,
+            module=r"earthaccess\.(results|store)",
+        )
+        yield
+
+
+def _login(credentials=None):
+    """Authenticate with earthaccess, handling credentials if provided.
+
+    Parameters
+    ----------
+    credentials : tuple(str, str) | None
+        (username, password) tuple, or None for anonymous login.
+    """
+    if credentials:
+        username, password = credentials
+        if username and password:
+            try:
+                earthaccess.login(strategy="environment", persist=True)
+            except Exception:
+                earthaccess.login()
+        else:
+            logger.warning(
+                "No Earthdata credentials provided, attempting anonymous login"
+            )
+            earthaccess.login()
+    else:
+        earthaccess.login()
+
+
+def _search(**params):
+    """Search earthaccess for granules matching query parameters.
+
+    Parameters
+    ----------
+    **params
+        Query parameters (short_name, temporal, bounding_box, granule_name, etc.).
+
+    Returns
+    -------
+    list
+        Matching granule results, or empty list on error.
+    """
+    try:
+        with _suppress_granule_size_warning():
+            results = earthaccess.search_data(**params)
+        return results
+    except Exception as e:
+        logger.error("Error searching for data: %s", e)
+        return []
+
+
+def _filter_new(results, processed_granules):
+    """Filter results to keep only granules not yet processed.
+
+    Parameters
+    ----------
+    results : list
+        Granule result objects from earthaccess.search_data().
+    processed_granules : set[str]
+        GranuleUR values already processed, to exclude.
+
+    Returns
+    -------
+    list
+        Granule result objects not in processed_granules.
+    """
+    granule_ids = [result["umm"]["GranuleUR"] for result in results]
+    new_granule_ids = [gid for gid in granule_ids if gid not in processed_granules]
+    return [r for r in results if r["umm"]["GranuleUR"] in new_granule_ids]
+
+
+def _download_files(results, directory):
+    """Download granule files to a local directory.
+
+    Parameters
+    ----------
+    results : list
+        Granule result objects to download.
+    directory : str
+        Local directory to write files to.
+
+    Returns
+    -------
+    list
+        Paths to successfully downloaded files, or empty list on error.
+    """
+    if not results:
+        return []
+
+    try:
+        with _suppress_granule_size_warning():
+            files = earthaccess.download(results, directory, show_progress=True)
+        return files or []
+    except Exception as e:
+        logger.error("Error downloading files: %s", e)
+        return []
+
+
 def query(
     aoi: list,
     startdate: datetime.date,
     enddate: datetime.date,
     product: str = SWOT_LAKE_SHORT_NAME,
-    earthaccess_client=earthaccess,
 ) -> object:
     # format coordinates and extract bounds
     aoi = geometry.format_coord_list(aoi)
 
-    # login and authenticate earthacess
-    earthaccess_client.login()
+    # login and authenticate earthaccess
+    _login()
 
     # define query parameters
     params = {
@@ -35,21 +141,11 @@ def query(
         "bounding_box": shapely.Polygon(aoi).bounds,
     }
 
-    # Silence a known earthaccess deprecation warning until upstream migrates
-    # DataGranule.size() to DataGranule.size attribute access.
-    with warnings.catch_warnings():
-        warnings.filterwarnings(
-            "ignore",
-            message=r"As of version 1\.0, `DataGranule\.size` will be accessed as an attribute",
-            category=FutureWarning,
-            module=r"earthaccess\.results",
-        )
-        results = earthaccess_client.search_data(**params)
-
+    results = _search(**params)
     return results
 
 
-def download(results, download_directory: str, earthaccess_client=earthaccess):
+def download(results, download_directory: str):
     # Check if we have a progress log file in this directory, if not make it
     log_path = os.path.join(download_directory, "downloaded.log")
     if not os.path.exists(log_path):
@@ -71,16 +167,7 @@ def download(results, download_directory: str, earthaccess_client=earthaccess):
     logger.info("%s files shown as downloaded in log", len(results) - len(to_download))
     logger.info("%s files will be downloaded", len(to_download))
     if to_download:
-        with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore",
-                message=r"As of version 1\.0, `DataGranule\.size` will be accessed as an attribute",
-                category=FutureWarning,
-                module=r"earthaccess\.(results|store)",
-            )
-            files = earthaccess_client.download(
-                to_download, download_directory, show_progress=True
-            )
+        files = _download_files(to_download, download_directory)
 
         with open(log_path, "a") as log:
             for file in files:

HydroEO/satellites/swot/pixc.py

@@ -11,7 +11,6 @@
 from pathlib import Path
 from typing import Any
 
-import earthaccess
 import geopandas as gpd
 import numpy as np
 import pandas as pd
@@ -23,6 +22,13 @@
 from shapely.geometry import box
 from tqdm import tqdm
 
+from HydroEO.satellites.swot.download import (
+    _download_files,
+    _filter_new,
+    _login,
+    _search,
+)
+
 logger = logging.getLogger(__name__)
 
 # Classification enum for SWOT PIXC data
@@ -116,15 +122,7 @@ def _download_granules(
     """Download SWOT PIXC granules matching query, skipping already processed ones."""
     logger.debug("=== SWOT Pixel Cloud Download Phase ===")
 
-    username, password = credentials
-    if username and password:
-        try:
-            earthaccess.login(strategy="environment", persist=True)
-        except Exception:
-            earthaccess.login()
-    else:
-        logger.warning("No Earthdata credentials provided, attempting anonymous login")
-        earthaccess.login()
+    _login(credentials)
 
     aoi_config = config["aoi"]
     if aoi_config["type"] == "bbox":
@@ -147,44 +145,34 @@ def _download_granules(
     logger.info(
         "Searching for SWOT PIXC product '%s' in AOI %s", product, aoi_config["name"]
     )
-    try:
-        swot_results = earthaccess.search_data(
-            short_name=product,
-            bounding_box=bounds,
-            temporal=(time_start, time_end),
-        )
-        logger.debug("Found %d granules matching query", len(swot_results))
-    except Exception as e:
-        logger.error("Error searching for SWOT data: %s", e)
-        return []
+
+    search_params = {
+        "short_name": product,
+        "bounding_box": bounds,
+        "temporal": (time_start, time_end),
+    }
+    swot_results = _search(**search_params)
+    logger.debug("Found %d granules matching query", len(swot_results))
 
     if not swot_results:
         logger.warning("No SWOT PIXC granules found matching the query")
         return []
 
-    granule_ids = [result["umm"]["GranuleUR"] for result in swot_results]
-    new_granules = [gid for gid in granule_ids if gid not in processed_granules]
+    new_results = _filter_new(swot_results, processed_granules)
     logger.debug(
         "%d new granules to download (already processed: %d)",
-        len(new_granules),
+        len(new_results),
         len(processed_granules),
     )
 
-    if not new_granules:
+    if not new_results:
         logger.info("All granules already processed, skipping download")
         return []
 
-    new_results = [r for r in swot_results if r["umm"]["GranuleUR"] in new_granules]
     logger.debug("Downloading %d granules to %s", len(new_results), raw_dir)
-    try:
-        downloaded_files = earthaccess.download(
-            new_results, str(raw_dir), show_progress=True
-        )
-        logger.debug("Successfully downloaded %d files", len(downloaded_files))
-        return downloaded_files or []
-    except Exception as e:
-        logger.error("Error downloading SWOT data: %s", e)
-        return []
+    downloaded_files = _download_files(new_results, str(raw_dir))
+    logger.debug("Successfully downloaded %d files", len(downloaded_files))
+    return downloaded_files or []
 
 
 def _preprocess_granules(

HydroEO/satellites/swot/raster.py

@@ -12,7 +12,6 @@
 from pathlib import Path
 from typing import Any
 
-import earthaccess
 import geopandas as gpd
 import numpy as np
 import rasterio
@@ -25,6 +24,12 @@
 from tqdm import tqdm
 
 from HydroEO import FLOAT32_NODATA_VALUE
+from HydroEO.satellites.swot.download import (
+    _download_files,
+    _filter_new,
+    _login,
+    _search,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -121,15 +126,7 @@ def _download_granules(
     """Download SWOT granules matching query, skipping already processed ones."""
     logger.debug("=== SWOT Raster Download Phase ===")
 
-    username, password = credentials
-    if username and password:
-        try:
-            earthaccess.login(strategy="environment", persist=True)
-        except Exception:
-            earthaccess.login()
-    else:
-        logger.warning("No Earthdata credentials provided, attempting anonymous login")
-        earthaccess.login()
+    _login(credentials)
 
     aoi_config = config["aoi"]
     if aoi_config["type"] == "bbox":
@@ -156,45 +153,35 @@ def _download_granules(
     logger.info(
         "Searching for SWOT product '%s' in AOI %s", product, aoi_config["name"]
     )
-    try:
-        swot_results = earthaccess.search_data(
-            short_name=product,
-            bounding_box=bounds,
-            temporal=(time_start, time_end),
-            granule_name=granule_filter,
-        )
-        logger.debug("Found %d granules matching query", len(swot_results))
-    except Exception as e:
-        logger.error("Error searching for SWOT data: %s", e)
-        return []
+
+    search_params = {
+        "short_name": product,
+        "bounding_box": bounds,
+        "temporal": (time_start, time_end),
+        "granule_name": granule_filter,
+    }
+    swot_results = _search(**search_params)
+    logger.debug("Found %d granules matching query", len(swot_results))
 
     if not swot_results:
         logger.warning("No SWOT granules found matching the query")
         return []
 
-    granule_ids = [result["umm"]["GranuleUR"] for result in swot_results]
-    new_granules = [gid for gid in granule_ids if gid not in processed_granules]
+    new_results = _filter_new(swot_results, processed_granules)
     logger.debug(
         "%d new granules to download (already processed: %d)",
-        len(new_granules),
+        len(new_results),
         len(processed_granules),
     )
 
-    if not new_granules:
+    if not new_results:
         logger.info("All granules already processed, skipping download")
         return []
 
-    new_results = [r for r in swot_results if r["umm"]["GranuleUR"] in new_granules]
     logger.debug("Downloading %d granules to %s", len(new_results), raw_dir)
-    try:
-        downloaded_files = earthaccess.download(
-            new_results, str(raw_dir), show_progress=True
-        )
-        logger.debug("Successfully downloaded %d files", len(downloaded_files))
-        return downloaded_files or []
-    except Exception as e:
-        logger.error("Error downloading SWOT data: %s", e)
-        return []
+    downloaded_files = _download_files(new_results, str(raw_dir))
+    logger.debug("Successfully downloaded %d files", len(downloaded_files))
+    return downloaded_files or []
 
 
 def _preprocess_granules(