Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions dataretrieval/waterdata/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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``
Expand All @@ -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:
Expand All @@ -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
)
Expand Down
8 changes: 8 additions & 0 deletions tests/waterdata_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down