diff --git a/dataretrieval/ngwmn.py b/dataretrieval/ngwmn.py index 154899fa..cc0cf9bf 100644 --- a/dataretrieval/ngwmn.py +++ b/dataretrieval/ngwmn.py @@ -105,6 +105,7 @@ def _get(service: str, local_vars: dict[str, Any]) -> tuple[pd.DataFrame, BaseMe service, output_id=_NGWMN_OUTPUT_ID, base_url=NGWMN_OGC_API_URL, + spatial=service == "sites", dialect=NGWMN_DIALECT, ) diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 5c5236d8..be63e69c 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -145,7 +145,10 @@ async def _paginate( def _ogc_parse_response( - resp: httpx.Response, *, geopd: bool + resp: httpx.Response, + *, + geopd: bool, + include_geometry: bool = True, ) -> tuple[pd.DataFrame, str | None]: """Parse one OGC API page: extract the DataFrame and the next-page URL. @@ -156,7 +159,12 @@ def _ogc_parse_response( """ body = resp.json() return ( - _get_resp_data(resp, geopd=geopd, body=body), + _get_resp_data( + resp, + geopd=geopd, + include_geometry=include_geometry, + body=body, + ), _next_req_url(resp, body=body) or None, ) @@ -166,6 +174,7 @@ async def _walk_pages( req: httpx.Request, client: httpx.AsyncClient | None = None, *, + include_geometry: bool = True, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: """ @@ -179,7 +188,10 @@ async def _walk_pages( Parameters ---------- geopd : bool - Whether geopandas is installed (drives geometry handling). + Whether this request's pages use an active GeoDataFrame. + include_geometry : bool, optional + Whether plain-DataFrame pages retain raw coordinate lists. False for + nonspatial collections and requests using ``skip_geometry=True``. req : httpx.Request The initial HTTP request to send. client : httpx.AsyncClient, optional @@ -216,7 +228,11 @@ async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: return await _paginate( req, - parse_response=functools.partial(_ogc_parse_response, geopd=geopd), + parse_response=functools.partial( + _ogc_parse_response, + geopd=geopd, + include_geometry=include_geometry, + ), follow_up=follow_up, client=client, row_cap=row_cap, @@ -229,6 +245,7 @@ def get_ogc_data( output_id: str, *, base_url: str, + spatial: bool, max_rows: int | None = None, extra_id_cols: frozenset[str] | set[str] = frozenset(), dialect: OgcDialect | None = None, @@ -266,6 +283,10 @@ def get_ogc_data( rejected at send time, surfacing as a NetworkError about an unknown service. Requiring it moves that mistake to the call site, where mypy catches it. + spatial : bool + Whether this collection's result contract includes feature geometry. + The adapter supplies this semantic fact; ``skip_geometry`` and + geopandas availability then select the concrete frame representation. extra_id_cols : set or frozenset, optional Synthetic id columns to push to the end of a result frame (see :func:`_arrange_cols`). Defaults to an empty set. @@ -314,6 +335,12 @@ def get_ogc_data( convert_type = args.pop("convert_type", False) args = {k: v for k, v in args.items() if v is not None} + # Choose one semantic frame shape for the whole request. Every page and + # the all-empty finalizer receive these same values, so frame type never + # depends on whether a particular page happens to carry geometry. + include_geometry = spatial and not bool(args.get("skip_geometry", False)) + geopd = GEOPANDAS and include_geometry + # Post-processing is injected into the chunker rather than applied here, # so it runs on *every* exit: the normal return AND a later # ``exc.call.resume()`` after a ChunkInterrupted (which never re-enters @@ -328,6 +355,8 @@ def get_ogc_data( output_id=output_id, convert_type=convert_type, collection=collection, + geopd=geopd, + include_geometry=include_geometry, max_rows=max_rows, extra_id_cols=extra_id_cols, dialect=dialect, @@ -354,7 +383,12 @@ def get_ogc_data( # ``(df, BaseMetadata)`` shape rather than a raw response pair. return FanOut( [req], - functools.partial(_walk_pages, GEOPANDAS, row_cap=max_rows), + functools.partial( + _walk_pages, + geopd, + include_geometry=include_geometry, + row_cap=max_rows, + ), RetryPolicy.from_env(), finalize, canonical_url=str(req.url), @@ -371,7 +405,11 @@ def get_ogc_data( _construct_api_requests, base_url=base_url, dialect=dialect ) fetch = functools.partial( - _fetch_once, build_request=build_request, row_cap=max_rows + _fetch_once, + build_request=build_request, + geopd=geopd, + include_geometry=include_geometry, + row_cap=max_rows, ) run = chunking.multi_value_chunked(build_request=build_request)(fetch) # No progress block here: the executor that emits the events owns the line @@ -383,6 +421,8 @@ async def _fetch_once( args: dict[str, Any], *, build_request: Callable[..., httpx.Request], + geopd: bool, + include_geometry: bool, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: """Send one prepared-args OGC request asynchronously; return (frame, response). @@ -400,4 +440,9 @@ async def _fetch_once( synchronously. The return shape is ``(frame, response)``. """ req = build_request(**args) - return await _walk_pages(geopd=GEOPANDAS, req=req, row_cap=row_cap) + return await _walk_pages( + geopd=geopd, + req=req, + include_geometry=include_geometry, + row_cap=row_cap, + ) diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index d97d58b5..760e6f41 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -43,16 +43,33 @@ ) -def _empty_feature_frame(geopd: bool) -> pd.DataFrame: - """Empty result frame for a page that carries no features. - - Returns a ``GeoDataFrame`` when geopandas is available so a downstream - ``pd.concat([empty_page, geo_page])`` doesn't downgrade a geopandas - user's result to a plain ``DataFrame`` (stripping geometry/CRS). The - single home for this empty-page contract, shared by the feature-frame - builders that flatten GeoJSON pages. +def _empty_feature_frame( + geopd: bool, + columns: list[str] | None = None, + *, + include_geometry: bool = True, +) -> pd.DataFrame: + """Construct an empty feature frame for one request's geometry mode. + + ``geopd`` and ``include_geometry`` are selected once before pagination and + reused for every page and the final all-empty result. This keeps empty and + non-empty pages in the same frame family, so ordinary ``pd.concat`` is + sufficient. ``columns`` is supplied only when finalization has fetched the + collection schema; page-level empties intentionally remain schema-light. """ - return gpd.GeoDataFrame() if geopd else pd.DataFrame() + result_columns = list(columns or []) + if include_geometry: + if "geometry" not in result_columns: + result_columns.append("geometry") + else: + result_columns = [name for name in result_columns if name != "geometry"] + + data = {name: pd.Series(dtype=object) for name in result_columns} + if not geopd: + return pd.DataFrame(data, columns=result_columns) + + data["geometry"] = gpd.GeoSeries([], crs=_CRS) + return gpd.GeoDataFrame(data, columns=result_columns, geometry="geometry", crs=_CRS) def _attach_coordinates(df: pd.DataFrame, features: list[dict[str, Any]]) -> None: @@ -86,6 +103,7 @@ def _get_resp_data( resp: httpx.Response, geopd: bool, *, + include_geometry: bool = True, body: dict[str, Any] | None = None, ) -> pd.DataFrame: """ @@ -97,8 +115,11 @@ def _get_resp_data( The HTTP response object expected to contain a JSON body with a "features" key. geopd : bool - Indicates whether geopandas is installed and should be used to - handle geometries. + Whether this request's result contract uses an active GeoDataFrame. + Selected once before pagination rather than inferred from page content. + include_geometry : bool, optional + Whether plain-DataFrame results retain raw coordinate lists. False for + nonspatial collections and ``skip_geometry=True`` requests. body : dict, optional Pre-parsed JSON body for ``resp``. When provided, skips the ``resp.json()`` call — useful when the caller has already @@ -132,11 +153,11 @@ def _get_resp_data( # actually carry features. An absent/empty ``features`` is also the real # schema-drift shape (a 200 with no features) — treat it as empty rather # than crash with a ``KeyError`` downstream, which ``_paginate`` would - # mistake for a transient transport error. ``_empty_feature_frame`` - # preserves the GeoDataFrame type on the short-circuit (see its docstring). + # mistake for a transient transport error. The request-level mode gives + # this page the same frame family as every non-empty page in the walk. features = body.get("features") or [] if not features: - return _empty_feature_frame(geopd) + return _empty_feature_frame(geopd, include_geometry=include_geometry) if not geopd: properties = [feature.get("properties") or {} for feature in features] @@ -144,25 +165,18 @@ def _get_resp_data( # Always materialize the feature-level ID (possibly all-None) so # ``_arrange_cols`` can perform the documented collection-specific rename. df["id"] = [feature.get("id") for feature in features] - _attach_coordinates(df, features) + if include_geometry: + _attach_coordinates(df, features) return df - # Organize json into geodataframe and make sure id column comes along. - # NGWMN observation collections (water levels, lithology, …) return - # features with no ``geometry`` key at all; ``_geo_feature_frame`` absorbs - # that, and the all-null check below then yields a plain DataFrame. + # A spatial request remains geospatial even when this particular page + # carries only null/missing geometries; changing frame family based on page + # contents makes pagination concat order-dependent. df = _geo_feature_frame(features) - # Mirror the non-geopandas branch's defensive ``f.get("id")`` so a feature - # missing a top-level ``id`` yields None rather than a KeyError. + # Mirror the plain branch's defensive ``f.get("id")`` so a feature missing + # a top-level id yields None rather than a KeyError. df["id"] = [f.get("id") for f in features] - df = df[["id"] + [col for col in df.columns if col != "id"]] - - # If no geometry present, then return pandas dataframe. A geodataframe - # is not needed. - if df["geometry"].isnull().all(): - df = pd.DataFrame(df.drop(columns="geometry")) - - return df + return df[["id"] + [col for col in df.columns if col != "id"]] def _deal_with_empty( @@ -171,44 +185,32 @@ def _deal_with_empty( collection: str, *, base_url: str, + geopd: bool, + include_geometry: bool, ) -> pd.DataFrame: + """Apply the collection schema when an entire result is empty. + + Page-level empties already carry the request's frame family. The complete + column list is available only from explicit ``properties`` or the + collection schema, so construct the fully shaped empty once here rather + than fetching schema during pagination. """ - Handles empty DataFrame results by returning a DataFrame with appropriate columns. + if not return_list.empty: + return return_list - If `return_list` is empty, determines the column names to use: - - If `properties` is not provided or contains only NaN values, - retrieves schema properties from the specified collection. - - Otherwise, uses the provided `properties` list as column names. + if not properties or all(pd.isna(properties)): + from dataretrieval.ogc.schema import _check_ogc_requests - Parameters - ---------- - return_list : pd.DataFrame - The DataFrame to check for emptiness. - properties : Optional[List[str]] - List of property names to use as columns, or None. - collection : str - The collection endpoint to query for schema properties if needed. - base_url : str - OGC API base URL to use for that schema query — the API the empty - result came from, not any particular collection. - - Returns - ------- - pd.DataFrame - The original DataFrame if not empty, otherwise an empty - DataFrame with the appropriate columns. - """ - if return_list.empty: - if not properties or all(pd.isna(properties)): - # Schema lookup performs HTTP only for an empty result. - from dataretrieval.ogc.schema import _check_ogc_requests + schema, _ = _check_ogc_requests( + endpoint=collection, req_type="schema", base_url=base_url + ) + properties = list(schema.get("properties", {}).keys()) - schema, _ = _check_ogc_requests( - endpoint=collection, req_type="schema", base_url=base_url - ) - properties = list(schema.get("properties", {}).keys()) - return pd.DataFrame(columns=properties) - return return_list + return _empty_feature_frame( + geopd, + properties, + include_geometry=include_geometry, + ) def _arrange_cols( @@ -361,6 +363,8 @@ def _finalize_ogc( output_id: str, convert_type: bool, collection: str, + geopd: bool, + include_geometry: bool, max_rows: int | None = None, extra_id_cols: frozenset[str] | set[str] = frozenset(), dialect: OgcDialect | None = None, @@ -389,7 +393,14 @@ def _finalize_ogc( """ if dialect is None: dialect = DEFAULT_DIALECT - frame = _deal_with_empty(frame, properties, collection, base_url=base_url) + frame = _deal_with_empty( + frame, + properties, + collection, + base_url=base_url, + geopd=geopd, + include_geometry=include_geometry, + ) # Normalize to PEP-8 snake_case column names *first*, so the dialect's # ``time_cols``/``numerical_cols``/``sort_cols`` (all snake_case) match # regardless of whether the API returns snake_case (Water Data, where diff --git a/dataretrieval/waterdata/reference.py b/dataretrieval/waterdata/reference.py index 14a1884b..1d988640 100644 --- a/dataretrieval/waterdata/reference.py +++ b/dataretrieval/waterdata/reference.py @@ -112,7 +112,11 @@ def get_reference_table( if limit is not None: query_args["limit"] = limit return get_ogc_data( - args=query_args, output_id=output_id, collection=collection, max_rows=max_rows + args=query_args, + output_id=output_id, + collection=collection, + max_rows=max_rows, + spatial=False, ) diff --git a/dataretrieval/waterdata/utils.py b/dataretrieval/waterdata/utils.py index d2cc9182..3a0ec5c3 100644 --- a/dataretrieval/waterdata/utils.py +++ b/dataretrieval/waterdata/utils.py @@ -178,6 +178,8 @@ def get_ogc_data( output_id: str | None = None, max_rows: int | None = None, cql_body: str | None = None, + *, + spatial: bool = True, ) -> tuple[pd.DataFrame, BaseMetadata]: """Water-Data wrapper over :func:`~dataretrieval.ogc.get_ogc_data`. @@ -204,6 +206,9 @@ def get_ogc_data( cql_body : str, optional A verbatim CQL2 JSON body to POST instead of building the query from ``args`` (see the facade's ``cql_body``). Used by :func:`get_cql`. + spatial : bool, optional + Whether the collection carries feature geometry. Water Data's typed + feature collections do; reference tables pass ``False``. Returns ------- @@ -221,6 +226,7 @@ def get_ogc_data( output_id, max_rows=max_rows, base_url=OGC_API_URL, + spatial=spatial, extra_id_cols=_EXTRA_ID_COLS, dialect=WATERDATA_DIALECT, cql_body=cql_body, diff --git a/tests/ngwmn_test.py b/tests/ngwmn_test.py index 5921aa1e..78dd0fac 100644 --- a/tests/ngwmn_test.py +++ b/tests/ngwmn_test.py @@ -307,6 +307,28 @@ def test_get_sites_skip_geometry(httpx_mock): assert _queries(httpx_mock, "sites")[0]["skipGeometry"] == ["true"] +def test_get_sites_empty_skip_geometry_is_plain(httpx_mock): + """The request's skip-geometry mode also governs an all-empty result.""" + httpx_mock.add_response( + method="GET", + url=_schema_re("sites"), + json={ + "properties": { + "geometry": {}, + "monitoring_location_id": {}, + "state_name": {}, + } + }, + ) + _mock(httpx_mock, "sites", _collection([])) + + df, _ = ngwmn.get_sites(monitoring_location_id="USGS-00000000", skip_geometry=True) + + assert type(df) is DataFrame + assert "geometry" not in df.columns + assert "monitoring_location_id" in df.columns + + def test_get_sites_state_accepts_name_postal_or_fips(httpx_mock): """The single ``state`` parameter accepts a full name, postal code, or FIPS code, and all three are normalized to the full ``state_name`` that the @@ -340,6 +362,28 @@ def test_get_providers(httpx_mock): assert "geometry" not in df.columns +def test_get_providers_empty_stays_plain(httpx_mock): + """A nonspatial collection stays tabular when its query has no matches.""" + httpx_mock.add_response( + method="GET", + url=_schema_re("providers"), + json={ + "properties": { + "agency_code": {}, + "organization_type": {}, + "state": {}, + } + }, + ) + _mock(httpx_mock, "providers", _collection([])) + + df, _ = ngwmn.get_providers(agency_code="NONE") + + assert type(df) is DataFrame + assert "geometry" not in df.columns + assert {"agency_code", "organization_type", "state"}.issubset(df.columns) + + def test_get_providers_state_accepts_name_postal_or_fips(httpx_mock): """``get_providers`` normalizes any state encoding to the uppercase postal code that the ``providers`` collection queries on -- the other half of the @@ -515,8 +559,10 @@ def test_empty_result_returns_typed_empty_frame(httpx_mock): df, _ = ngwmn.get_water_level(monitoring_location_id=_SITE) + assert type(df) is DataFrame assert df.empty assert "monitoring_location_id" in df.columns + assert "geometry" not in df.columns # --- live upstream monitor --------------------------------------------------- diff --git a/tests/waterdata_filters_test.py b/tests/waterdata_filters_test.py index 9ca439c6..ff6c6f25 100644 --- a/tests/waterdata_filters_test.py +++ b/tests/waterdata_filters_test.py @@ -164,7 +164,7 @@ def test_long_filter_fans_out_into_multiple_requests(): expr = _filter_chunking_clauses() sent_filters: list[str] = [] - async def fake_walk_pages(*, geopd, req, row_cap=None): + async def fake_walk_pages(*, geopd, req, include_geometry=True, row_cap=None): idx = len(sent_filters) sent_filters.append(_query_params(req).get("filter", [None])[0]) return pd.DataFrame({"id": [f"chunk-{idx}"], "value": [idx]}), _fake_response() diff --git a/tests/waterdata_test.py b/tests/waterdata_test.py index 7d8fe80f..e3436925 100644 --- a/tests/waterdata_test.py +++ b/tests/waterdata_test.py @@ -762,6 +762,29 @@ def test_get_daily_no_geometry(httpx_mock): assert _sent(httpx_mock, "daily")[0]["skipGeometry"] == ["true"] +def test_get_daily_empty_no_geometry(httpx_mock): + """An empty skip-geometry request has the same plain shape as a hit.""" + httpx_mock.add_response( + method="GET", + url=_schema_url("daily"), + json={ + "properties": { + "geometry": {}, + "id": {}, + "monitoring_location_id": {}, + "value": {}, + } + }, + ) + _mock_items(httpx_mock, "daily", body={"features": [], "links": []}) + + df, _ = get_daily(monitoring_location_id="USGS-NOT-FOUND", skip_geometry=True) + + assert type(df) is DataFrame + assert "geometry" not in df.columns + assert {"monitoring_location_id", "value"}.issubset(df.columns) + + def test_get_continuous(httpx_mock): """Continuous observations are timestamped (not date-only), so ``time`` comes back as a UTC-aware datetime column.""" diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index 9081b701..d76fb7d1 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -33,6 +33,7 @@ from dataretrieval.ogc.schema import _check_ogc_requests from dataretrieval.ogc.shaping import ( _arrange_cols, + _deal_with_empty, _get_resp_data, _to_snake_case, ) @@ -53,6 +54,8 @@ extra_id_cols=_EXTRA_ID_COLS, dialect=WATERDATA_DIALECT, base_url=OGC_API_URL, + geopd=False, + include_geometry=True, ) _LOGGER_NAME = _utils_module.__name__ @@ -605,6 +608,66 @@ class _Sentinel: assert isinstance(result, _Sentinel) +def test_empty_result_uses_request_geometry_contract(): + """All-empty shaping follows the request mode, not the template class.""" + with_geometry = _deal_with_empty( + pd.DataFrame(), + ["a", "b"], + "daily", + base_url="x", + geopd=False, + include_geometry=True, + ) + without_geometry = _deal_with_empty( + pd.DataFrame(), + ["a", "geometry", "b"], + "daily", + base_url="x", + geopd=False, + include_geometry=False, + ) + + assert type(with_geometry) is pd.DataFrame + assert list(with_geometry.columns) == ["a", "b", "geometry"] + assert all(dtype.kind == "O" for dtype in with_geometry.dtypes) + assert list(without_geometry.columns) == ["a", "b"] + + +@pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") +def test_empty_result_matches_a_non_empty_geodataframe(): + """An empty geospatial result remains usable like a non-empty one.""" + import geopandas as gpd + + response = _resp_ok([]) + page = _get_resp_data(response, geopd=True) + template = pd.concat([page], ignore_index=True) + out = _deal_with_empty( + template, + ["monitoring_location_id"], + "daily", + base_url="x", + geopd=True, + include_geometry=True, + ) + + assert isinstance(template, gpd.GeoDataFrame) + assert isinstance(out, gpd.GeoDataFrame) and out.empty + assert out["monitoring_location_id"].dtype == object + assert out.geometry is not None + assert out.crs == "EPSG:4326" + out["monitoring_location_id"].str.startswith("USGS") + out.to_crs("EPSG:3857") + real = gpd.GeoDataFrame( + { + "monitoring_location_id": ["USGS-1"], + "geometry": gpd.points_from_xy([1], [2]), + }, + crs="EPSG:4326", + ) + combined = pd.concat([out, real], ignore_index=True) + assert isinstance(combined, gpd.GeoDataFrame) and combined.crs == "EPSG:4326" + + def test_get_resp_data_empty_preserves_geopd_type(): """Same as the stats-side preservation: ``_get_resp_data``'s ``numberReturned == 0`` short-circuit must return a @@ -627,6 +690,27 @@ class _Sentinel: assert isinstance(result, _Sentinel) +@pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") +def test_get_resp_data_keeps_one_geospatial_type_when_geometry_is_missing(): + """A spatial request cannot change frame family based on page contents.""" + import geopandas as gpd + + empty = _get_resp_data(_resp_ok([]), geopd=True) + missing = _get_resp_data( + _resp_ok([{"id": "1", "properties": {"value": 1}}]), + geopd=True, + ) + + assert isinstance(empty, gpd.GeoDataFrame) + assert isinstance(missing, gpd.GeoDataFrame) + assert missing.geometry.isna().all() + for pages in ([empty, missing], [missing, empty]): + combined = pd.concat(pages, ignore_index=True) + assert isinstance(combined, gpd.GeoDataFrame) + assert combined.geometry.name == "geometry" + assert combined.crs == "EPSG:4326" + + def test_get_resp_data_attaches_wgs84_crs(): """A geometry-bearing Water Data page should come back tagged as EPSG:4326 (the CRS the coordinates are published in), so callers can