Skip to content

Commit

Permalink
usgs: Handle dataretrieval state codes
Browse files Browse the repository at this point in the history
dataretrieval changed the datatype of `state_codes` from a list to a dict.
With this commit we handle this change on searvey side.

Fixes #143
  • Loading branch information
pmav99 committed Jun 17, 2024
1 parent 28cb4b5 commit b9e4600
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions searvey/usgs.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# `dataretrieval` package developed by USGS is used for the USGS stations: https://usgs-python.github.io/dataretrieval/
#
# This pacakge is a thin wrapper around NWIS _REST API: https://waterservices.usgs.gov/rest/
# This package is a thin wrapper around NWIS _REST API: https://waterservices.usgs.gov/rest/
# We take the return values from `dataretrieval` to be the original data
from __future__ import annotations

Expand All @@ -22,7 +22,7 @@
import pandas as pd
import xarray as xr
from dataretrieval import nwis
from dataretrieval.codes import state_codes
from dataretrieval.codes import state_codes as _DATARETRIEVAL_STATE_CODES
from dataretrieval.utils import BaseMetadata
from shapely.geometry import MultiPolygon
from shapely.geometry import Polygon
Expand All @@ -43,6 +43,12 @@
# This will stop working if pandas switches its versioning scheme to CalVer or something...
_PANDAS_MAJOR_VERSION = int(importlib.metadata.version("pandas").split(".")[0])

# https://github.com/DOI-USGS/dataretrieval-python/compare/v1.0.8...v1.0.9
if isinstance(_DATARETRIEVAL_STATE_CODES, list):
STATE_CODES = sorted(_DATARETRIEVAL_STATE_CODES)
else:
STATE_CODES = sorted(_DATARETRIEVAL_STATE_CODES.values())

# constants
USGS_OUTPUT_OF_INTEREST = (
"elevation", # Overlaps some of the specific codes below
Expand Down Expand Up @@ -156,7 +162,7 @@ def _get_all_usgs_stations(normalize: bool = True) -> gpd.GeoDataFrame:
"hasDataType": dtp,
}
for st, dtp in product(
state_codes,
STATE_CODES,
USGS_OUTPUT_TYPE,
)
],
Expand Down

0 comments on commit b9e4600

Please sign in to comment.