Allow users to specify query parameters for the collections API when calling the get_reference_table function.

Author: jessamynsmith

dataretrieval/waterdata/api.py

@@ -1421,6 +1421,7 @@ def get_field_measurements(
 def get_reference_table(
         collection: str,
         limit: Optional[int] = None,
+        query: Optional[dict] = {},
         ) -> Tuple[pd.DataFrame, BaseMetadata]:
     """Get metadata reference tables for the USGS Water Data API.
 
@@ -1442,7 +1443,10 @@ def get_reference_table(
         allowable limit is 50000. It may be beneficial to set this number lower
         if your internet connection is spotty. The default (None) will set the
         limit to the maximum allowable limit for the service.
-    
+    query: dictionary, optional
+        The optional args parameter can be used to pass a dictionary of
+        query parameters to the collection API call.
+
     Returns
     -------
     df : ``pandas.DataFrame`` or ``geopandas.GeoDataFrame``
@@ -1463,6 +1467,12 @@ def get_reference_table(
         >>> ref, md = dataretrieval.waterdata.get_reference_table(
         ...     collection="parameter-codes"
         ... )
+
+        >>> # Get table of selected USGS parameter codes
+        >>> ref, md = dataretrieval.waterdata.get_reference_table(
+        ...     collection="parameter-codes"
+        ...     query={'id': '00001,00002'}
+        ... )
     """
     valid_code_services = get_args(METADATA_COLLECTIONS)
     if collection not in valid_code_services:

tests/waterdata_test.py

@@ -253,6 +253,14 @@ def test_get_reference_table():
     assert hasattr(md, 'url')
     assert hasattr(md, 'query_time')
 
+def test_get_reference_table_with_query():
+    query = {"id": "AK001,AK008"}
+    df, md = get_reference_table("agency-codes", query=query)
+    assert "agency_code" in df.columns
+    assert df.shape[0] == 2
+    assert hasattr(md, 'url')
+    assert hasattr(md, 'query_time')
+
 def test_get_reference_table_wrong_name():
     with pytest.raises(ValueError):
         get_reference_table("agency-cod")