From 3c0b35d0deccee6701b4142a1b5ce7c710f929a2 Mon Sep 17 00:00:00 2001 From: jessamynsmith Date: Fri, 13 Feb 2026 11:14:41 -0500 Subject: [PATCH] Allow users to specify query parameters for the collections API when calling the get_reference_table function. --- dataretrieval/waterdata/api.py | 14 ++++++++++++-- tests/waterdata_test.py | 8 ++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/dataretrieval/waterdata/api.py b/dataretrieval/waterdata/api.py index 4348d6a..4c1ee7e 100644 --- a/dataretrieval/waterdata/api.py +++ b/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: @@ -1480,7 +1490,7 @@ def get_reference_table( output_id = f"{collection.replace('-', '_')}" return get_ogc_data( - args={}, + args=query, output_id=output_id, service=collection ) diff --git a/tests/waterdata_test.py b/tests/waterdata_test.py index 3093408..89f2038 100755 --- a/tests/waterdata_test.py +++ b/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")