Merge pull request #19 from axiom-data-science/constraints

[OSRA-120] Allow passing search params as constraints to ERDDAP Source Objects

intake_erddap/erddap_cat.py

@@ -1,6 +1,6 @@
 """Catalog implementation for intake-erddap."""
 
-from typing import Dict, Optional, Tuple, Type, Union
+from typing import Dict, MutableMapping, Optional, Tuple, Type, Union
 
 import pandas as pd
 
@@ -33,9 +33,10 @@ class ERDDAPCatalog(Catalog):
     def __init__(
         self,
         server: str,
-        kwargs_search: Dict[str, Union[str, int, float]] = None,
+        kwargs_search: MutableMapping[str, Union[str, int, float]] = None,
         category_search: Optional[Tuple[str, str]] = None,
         erddap_client: Optional[Type[ERDDAP]] = None,
+        use_source_constraints: bool = True,
         **kwargs,
     ):
         """ERDDAPCatalog initialization
@@ -57,8 +58,14 @@ def __init__(
             custom_criteria key to narrow the search by, which will be matched to the category results
             using the custom_criteria that must be set up or input by the user, with `cf-pandas`.
             Currently only a single key can be matched at a time.
+        use_source_constraints : bool, default True
+            Any relevant search parameter defined in kwargs_search will be
+            passed to the source objects as contraints.
+
         """
         self._erddap_client = erddap_client or ERDDAP
+        self._entries: Dict[str, LocalCatalogEntry] = {}
+        self._use_source_contraints = use_source_constraints
         self.server = server
         self.search_url = None
 
@@ -131,7 +138,12 @@ def _load(self):
                 "server": self.server,
                 "dataset_id": dataset_id,
                 "protocol": "tabledap",
+                "constraints": {},
             }
+            if self._use_source_contraints and "min_time" in self.kwargs_search:
+                args["constraints"]["time>="] = self.kwargs_search["min_time"]
+            if self._use_source_contraints and "max_time" in self.kwargs_search:
+                args["constraints"]["time<="] = self.kwargs_search["max_time"]
 
             entry = LocalCatalogEntry(
                 dataset_id,

tests/test_erddap_cat.py

@@ -13,6 +13,14 @@
 from intake_erddap.erddap_cat import ERDDAPCatalog
 
 
+@pytest.fixture
+def single_dataset_catalog() -> pd.DataFrame:
+    """Fixture returns a dataframe with a single dataset ID."""
+    df = pd.DataFrame()
+    df["datasetID"] = ["abc123"]
+    return df
+
+
 def test_nothing():
     """This test exists to ensure that at least one test works."""
     pass
@@ -149,7 +157,7 @@ def test_params_search(mock_read_csv):
     df["datasetID"] = ["allDatasets", "abc123"]
     mock_read_csv.return_value = df
     erddap_url = "https://erddap.invalid/erddap"
-    search = search = {
+    search = {
         "min_lon": -100,
         "max_lon": -54,
         "min_lat": 19,
@@ -168,3 +176,23 @@ def test_params_search(mock_read_csv):
     assert query["minLon"] == "-100"
     assert int(float(query["minTime"])) == 1640995200
     assert query["standard_name"] == "sea_water_temperature"
+
+
+@mock.patch("pandas.read_csv")
+def test_constraints_present_in_source(mock_read_csv, single_dataset_catalog):
+    mock_read_csv.return_value = single_dataset_catalog
+    server = "https://erddap.invalid/erddap"
+    search = {
+        "min_time": "2022-01-01",
+        "max_time": "2022-11-07",
+    }
+    cat = ERDDAPCatalog(server=server, kwargs_search=search)
+    source = next(cat.values())
+    assert source._constraints["time>="] == "2022-01-01"
+    assert source._constraints["time<="] == "2022-11-07"
+
+    cat = ERDDAPCatalog(
+        server=server, kwargs_search=search, use_source_constraints=False
+    )
+    source = next(cat.values())
+    assert len(source._constraints) == 0