diff --git a/dataretrieval/exceptions.py b/dataretrieval/exceptions.py new file mode 100644 index 0000000..863277c --- /dev/null +++ b/dataretrieval/exceptions.py @@ -0,0 +1,6 @@ +class DataRetrievalError(BaseException): + """Base exception for dataretrieval""" + + +class EmptyQueryResultError(DataRetrievalError): + """Raised when a query returns an error""" diff --git a/dataretrieval/nwis.py b/dataretrieval/nwis.py index dfeb875..9c81626 100644 --- a/dataretrieval/nwis.py +++ b/dataretrieval/nwis.py @@ -11,6 +11,7 @@ from io import StringIO from dataretrieval.utils import to_str, format_datetime, update_merge, set_metadata as set_md +from .exceptions import EmptyQueryResultError from .utils import query WATERDATA_BASE_URL = 'https://nwis.waterdata.usgs.gov/' @@ -709,6 +710,10 @@ def _read_rdb(rdb): fields = rdb.splitlines()[count].split('\t') dtypes = {'site_no': str, 'dec_long_va': float, 'dec_lat_va': float} + # raise error if result has error + if "Error report" in (fields[0]): + raise EmptyQueryResultError + df = pd.read_csv(StringIO(rdb), delimiter='\t', skiprows=count + 2, names=fields, na_values='NaN', dtype=dtypes) diff --git a/tests/nwis_test.py b/tests/nwis_test.py index 22ecfca..e5a7541 100644 --- a/tests/nwis_test.py +++ b/tests/nwis_test.py @@ -1,5 +1,7 @@ import pytest -from dataretrieval.nwis import get_record + +from dataretrieval.exceptions import EmptyQueryResultError +from dataretrieval.nwis import get_record, _read_rdb START_DATE = '2018-01-24' END_DATE = '2018-01-25' @@ -39,6 +41,11 @@ def test_iv_service_answer(): assert df.index.names == [SITENO_COL, DATETIME_COL], "iv service returned incorrect index: {}".format(df.index.names) +def test_read_rdb_raises_error(): + with pytest.raises(EmptyQueryResultError): + _read_rdb("Error report") + + if __name__=='__main__': test_measurements_service_answer() test_iv_service_answer()