diff --git a/dataretrieval/streamstats.py b/dataretrieval/streamstats.py index 7cddabaa..9ae9339c 100644 --- a/dataretrieval/streamstats.py +++ b/dataretrieval/streamstats.py @@ -5,8 +5,6 @@ """ -import json - import requests @@ -58,7 +56,7 @@ def get_sample_watershed(): from the streamstats JSON object. """ - return get_watershed("NY", -74.524, 43.939) + return get_watershed("NY", -74.524, 43.939, format="watershed") def get_watershed( @@ -72,14 +70,13 @@ def get_watershed( simplify=True, format="geojson", ): - """Get watershed object based on location + """Delineate a watershed via the StreamStats API. - **Streamstats documentation:** - Returns a watershed object. The request configuration will determine the - overall request response. However all returns will return a watershed - object with at least the workspaceid. The workspace id is the id to the - service workspace where files are stored and can be used for further - processing such as for downloads and flow statistic computations. + Hits the StreamStats ``watershed.geojson`` endpoint and returns the + delineated watershed in one of three shapes selected by ``format``: the + raw ``requests.Response`` (default), the parsed JSON ``dict``, or a + :obj:`Watershed` instance. Every response carries a workspace identifier + that can be passed to :obj:`download_workspace` for further processing. See: https://streamstats.usgs.gov/streamstatsservices/#/ for more information. @@ -105,12 +102,17 @@ def get_watershed( simplify: bool, optional Boolean flag controlling whether or not to simplify the returned result. + format: string, optional + Selects the return shape. ``"geojson"`` (default) returns the raw + ``requests.Response``; ``"object"`` returns the parsed JSON ``dict``; + ``"watershed"`` returns a :obj:`Watershed` instance built from the + parsed JSON. Any other value raises ``ValueError``. Returns ------- - Watershed: :obj:`dataretrieval.streamstats.Watershed` - Custom object that contains the watershed information as extracted - from the streamstats JSON object. + requests.Response, dict, or :obj:`dataretrieval.streamstats.Watershed` + Watershed information from StreamStats. The exact return type + depends on ``format`` (see above). """ payload = { @@ -132,30 +134,52 @@ def get_watershed( if format == "geojson": return r - if format == "shape": - # use Fiona to return a shape object - pass + data = r.json() if format == "object": - # return a python object - pass + return data + + if format == "watershed": + return Watershed.from_streamstats_json(data) - data = json.loads(r.text) - return Watershed.from_streamstats_json(data) + raise ValueError( + f"Invalid format {format!r}; expected 'geojson', 'object', or 'watershed'." + ) class Watershed: - """Class to extract information from the streamstats JSON object.""" + """Class to extract information from the streamstats JSON object. - @classmethod - def from_streamstats_json(cls, streamstats_json): - """Method that creates a Watershed object from a streamstats JSON.""" - cls.watershed_point = streamstats_json["featurecollection"][0]["feature"] - cls.watershed_polygon = streamstats_json["featurecollection"][1]["feature"] - cls.parameters = streamstats_json["parameters"] - cls._workspaceID = streamstats_json["workspaceID"] - return cls + Attributes + ---------- + watershed_point : dict + GeoJSON feature for the watershed pour point. + watershed_polygon : dict + GeoJSON feature for the delineated watershed polygon. + parameters : list + Watershed parameters returned by StreamStats. + workspace_id : str + StreamStats workspace identifier for the watershed. + """ def __init__(self, rcode, xlocation, ylocation): - """Init method that calls the :obj:`from_streamstats_json` method.""" - get_watershed(rcode, xlocation, ylocation) + """Delineate a watershed and populate the instance from the response.""" + data = get_watershed(rcode, xlocation, ylocation, format="object") + self._populate_from_json(data) + + @classmethod + def from_streamstats_json(cls, streamstats_json): + """Construct a :obj:`Watershed` from a StreamStats JSON response.""" + instance = cls.__new__(cls) + instance._populate_from_json(streamstats_json) + return instance + + def _populate_from_json(self, streamstats_json): + self.watershed_point = streamstats_json["featurecollection"][0]["feature"] + self.watershed_polygon = streamstats_json["featurecollection"][1]["feature"] + self.parameters = streamstats_json["parameters"] + self.workspace_id = streamstats_json["workspaceID"] + + @property + def _workspaceID(self): + return self.workspace_id diff --git a/tests/streamstats_test.py b/tests/streamstats_test.py new file mode 100644 index 00000000..31a389f5 --- /dev/null +++ b/tests/streamstats_test.py @@ -0,0 +1,33 @@ +"""Tests for the streamstats module.""" + +from dataretrieval.streamstats import Watershed + +SAMPLE_JSON = { + "featurecollection": [ + {"name": "globalwatershedpoint", "feature": {"type": "Feature", "id": "pt-1"}}, + {"name": "globalwatershed", "feature": {"type": "Feature", "id": "poly-1"}}, + ], + "parameters": [{"code": "DRNAREA", "value": 41.2}], + "workspaceID": "NY20240101000000000", +} + + +def test_from_streamstats_json_does_not_mutate_class(): + """Two watersheds must not share state via class-level attributes.""" + other_json = { + "featurecollection": [ + {"feature": {"id": "pt-2"}}, + {"feature": {"id": "poly-2"}}, + ], + "parameters": [{"code": "OTHER", "value": 1.0}], + "workspaceID": "VT20240101000000000", + } + w1 = Watershed.from_streamstats_json(SAMPLE_JSON) + w2 = Watershed.from_streamstats_json(other_json) + + assert w1.workspace_id == "NY20240101000000000" + assert w2.workspace_id == "VT20240101000000000" + assert w1.parameters[0]["code"] == "DRNAREA" + assert w2.parameters[0]["code"] == "OTHER" + assert w1.watershed_point["id"] == "pt-1" + assert w2.watershed_point["id"] == "pt-2"