Skip to content

Commit b3a5d29

Browse files
thodson-usgsclaude
andcommitted
Make get_watershed always return a Watershed
Per Copilot review: the function's name and narrative docstring both promise a Watershed, but the actual contract was three different return types selected by `format` (Response, dict, Watershed) — and the default returned a `requests.Response`, contradicting the narrative. Drop the `format` parameter and always return a `Watershed`. Extract the HTTP+parse into a private `_fetch_streamstats_json` helper that both `get_watershed` and `Watershed.__init__` use, so they share fetch logic without circular calls. The narrative docstring now matches the contract. Also drop dead commented-out lines in `download_workspace` while in the file. Breaking change for callers using `format="geojson"`/`"object"`/etc. Callers needing the raw JSON or Response can use `Watershed.from_streamstats_json(...)` and `requests.get(...)` directly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent f3d1e8f commit b3a5d29

1 file changed

Lines changed: 48 additions & 54 deletions

File tree

dataretrieval/streamstats.py

Lines changed: 48 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,6 @@ def download_workspace(workspaceID, format=""):
3434

3535
r.raise_for_status()
3636
return r
37-
# data = r.raw.read()
38-
39-
# with open(filepath, 'wb') as f:
40-
# f.write(data)
41-
42-
# return
4337

4438

4539
def get_sample_watershed():
@@ -56,7 +50,35 @@ def get_sample_watershed():
5650
from the streamstats JSON object.
5751
5852
"""
59-
return get_watershed("NY", -74.524, 43.939, format="watershed")
53+
return get_watershed("NY", -74.524, 43.939)
54+
55+
56+
def _fetch_streamstats_json(
57+
rcode,
58+
xlocation,
59+
ylocation,
60+
crs=4326,
61+
includeparameters=True,
62+
includeflowtypes=False,
63+
includefeatures=True,
64+
simplify=True,
65+
):
66+
"""Hit the StreamStats watershed endpoint and return the parsed JSON."""
67+
payload = {
68+
"rcode": rcode,
69+
"xlocation": xlocation,
70+
"ylocation": ylocation,
71+
"crs": crs,
72+
"includeparameters": includeparameters,
73+
"includeflowtypes": includeflowtypes,
74+
"includefeatures": includefeatures,
75+
"simplify": simplify,
76+
}
77+
url = "https://streamstats.usgs.gov/streamstatsservices/watershed.geojson"
78+
79+
r = requests.get(url, params=payload)
80+
r.raise_for_status()
81+
return r.json()
6082

6183

6284
def get_watershed(
@@ -68,16 +90,14 @@ def get_watershed(
6890
includeflowtypes=False,
6991
includefeatures=True,
7092
simplify=True,
71-
format="geojson",
7293
):
73-
"""Get watershed object based on location
94+
"""Get watershed object based on location.
7495
75-
**Streamstats documentation:**
76-
Returns a watershed object. The request configuration will determine the
77-
overall request response. However all returns will return a watershed
78-
object with at least the workspaceid. The workspace id is the id to the
79-
service workspace where files are stored and can be used for further
80-
processing such as for downloads and flow statistic computations.
96+
Delineates a watershed via the StreamStats ``watershed.geojson`` endpoint
97+
and returns it as a :obj:`Watershed` instance carrying the pour point,
98+
polygon, parameters, and workspace identifier from the response. The
99+
workspace identifier can be passed to :obj:`download_workspace` for
100+
further processing.
81101
82102
See: https://streamstats.usgs.gov/streamstatsservices/#/ for more
83103
information.
@@ -103,49 +123,24 @@ def get_watershed(
103123
simplify: bool, optional
104124
Boolean flag controlling whether or not to simplify the returned
105125
result.
106-
format: string, optional
107-
Selects the return shape. ``"geojson"`` (default) returns the raw
108-
``requests.Response``; ``"object"`` returns the parsed JSON ``dict``;
109-
``"watershed"`` returns a :obj:`Watershed` instance built from the
110-
parsed JSON. Any other value raises ``ValueError``.
111126
112127
Returns
113128
-------
114-
requests.Response, dict, or :obj:`dataretrieval.streamstats.Watershed`
115-
Watershed information from StreamStats. The exact return type
116-
depends on ``format`` (see above).
129+
Watershed: :obj:`dataretrieval.streamstats.Watershed`
130+
Watershed instance built from the StreamStats response.
117131
118132
"""
119-
payload = {
120-
"rcode": rcode,
121-
"xlocation": xlocation,
122-
"ylocation": ylocation,
123-
"crs": crs,
124-
"includeparameters": includeparameters,
125-
"includeflowtypes": includeflowtypes,
126-
"includefeatures": includefeatures,
127-
"simplify": simplify,
128-
}
129-
url = "https://streamstats.usgs.gov/streamstatsservices/watershed.geojson"
130-
131-
r = requests.get(url, params=payload)
132-
133-
r.raise_for_status()
134-
135-
if format == "geojson":
136-
return r
137-
138-
data = r.json()
139-
140-
if format == "object":
141-
return data
142-
143-
if format == "watershed":
144-
return Watershed.from_streamstats_json(data)
145-
146-
raise ValueError(
147-
f"Invalid format {format!r}; expected 'geojson', 'object', or 'watershed'."
133+
data = _fetch_streamstats_json(
134+
rcode,
135+
xlocation,
136+
ylocation,
137+
crs=crs,
138+
includeparameters=includeparameters,
139+
includeflowtypes=includeflowtypes,
140+
includefeatures=includefeatures,
141+
simplify=simplify,
148142
)
143+
return Watershed.from_streamstats_json(data)
149144

150145

151146
class Watershed:
@@ -165,8 +160,7 @@ class Watershed:
165160

166161
def __init__(self, rcode, xlocation, ylocation):
167162
"""Delineate a watershed and populate the instance from the response."""
168-
data = get_watershed(rcode, xlocation, ylocation, format="object")
169-
self._populate_from_json(data)
163+
self._populate_from_json(_fetch_streamstats_json(rcode, xlocation, ylocation))
170164

171165
@classmethod
172166
def from_streamstats_json(cls, streamstats_json):

0 commit comments

Comments
 (0)