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

Author: jessamynsmith

dataretrieval/waterdata/api.py

@@ -234,7 +234,7 @@ def get_continuous(
     with the continuous endpoint. If the "time" input is left blank, the service
     will return the most recent year of measurements. Users may request no more
     than three years of data with each function call.
-    
+
     Continuous data are collected at a high frequency, typically 15-minute
     intervals. Depending on the specific monitoring location, the data may be
     transmitted automatically via telemetry and be available on WDFN within
@@ -243,7 +243,7 @@ def get_continuous(
     transmit data.  Continuous data are described by parameter name and
     parameter code (pcode). These data might also be referred to as
     "instantaneous values" or "IV".
-    
+
     Parameters
     ----------
     monitoring_location_id : string or list of strings, optional
@@ -482,11 +482,11 @@ def get_monitoring_locations(
         is located.
     county_code : string or list of strings, optional
         The code for the county or county equivalent (parish, borough, etc.) in which
-        the monitoring location is located. A `list of codes 
+        the monitoring location is located. A `list of codes
         <https://help.waterdata.usgs.gov/code/county_query?fmt=html>`_ is available.
     county_name : string or list of strings, optional
         The name of the county or county equivalent (parish, borough, etc.) in which
-        the monitoring location is located. A `list of codes 
+        the monitoring location is located. A `list of codes
         <https://help.waterdata.usgs.gov/code/county_query?fmt=html>`_ is available.
     minor_civil_division_code : string or list of strings, optional
         Codes for primary governmental or administrative divisions of the county or
@@ -1421,11 +1421,12 @@ 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.
 
     Reference tables provide the range of allowable values for parameter
-    arguments in the waterdata module. 
+    arguments in the waterdata module.
 
     Parameters
     ----------
@@ -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``
@@ -1454,7 +1458,7 @@ def get_reference_table(
         medium code values).
     md: :obj:`dataretrieval.utils.Metadata`
         A custom metadata object including the URL request and query time.
-    
+
     Examples
     --------
     .. code::
@@ -1463,24 +1467,30 @@ 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:
         raise ValueError(
             f"Invalid code service: '{collection}'. "
             f"Valid options are: {valid_code_services}."
         )
-    
+
     # Give ID column the collection name with underscores
     if collection.endswith("s") and collection != "counties":
         output_id = f"{collection[:-1].replace('-', '_')}"
     elif collection == "counties":
         output_id = "county"
     else:
         output_id = f"{collection.replace('-', '_')}"
-    
+
     return get_ogc_data(
-        args={},
+        args=query,
         output_id=output_id,
         service=collection
         )

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")