diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index f4a27f07..26131f96 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -20,7 +20,8 @@ jobs: python-version: "3.14" cache: "pip" - name: Install ruff - run: pip install ruff + # Keep this version aligned with the ruff-pre-commit revision. + run: pip install ruff==0.16.1 - name: Lint with ruff run: | ruff check . --output-format=github diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0813f1fe..8dd44167 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: - id: debug-statements - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.15 + rev: v0.16.1 hooks: - id: ruff-check args: [--fix] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2fbd526b..10b1ef4c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -149,19 +149,20 @@ via any automated processes or pipelines. * Example: ``` python + LIGHT_MESSAGES = { + "English": "There are %(number_of_lights)s lights.", + "Pirate": "Arr! Thar be %(number_of_lights)s lights.", + } - LIGHT_MESSAGES = { - 'English': "There are %(number_of_lights)s lights.", - 'Pirate': "Arr! Thar be %(number_of_lights)s lights." - } - def lights_message(language, number_of_lights): - """Return a language-appropriate string reporting the light count.""" - return LIGHT_MESSAGES[language] % locals() + def lights_message(language, number_of_lights): + """Return a language-appropriate string reporting the light count.""" + return LIGHT_MESSAGES[language] % locals() - def is_pirate(message): - """Return True if the given message sounds piratical.""" - return re.search(r"(?i)(arr|avast|yohoho)!", message) is not None + + def is_pirate(message): + """Return True if the given message sounds piratical.""" + return re.search(r"(?i)(arr|avast|yohoho)!", message) is not None ``` --- diff --git a/README.md b/README.md index a3f494fe..88b076ab 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,7 @@ and set it as an environment variable: ```python import os + os.environ["API_USGS_PAT"] = "your_api_key_here" ``` @@ -59,9 +60,9 @@ from dataretrieval import waterdata # Get daily streamflow data (returns DataFrame and metadata) df, metadata = waterdata.get_daily( - monitoring_location_id='USGS-01646500', - parameter_code='00060', # Discharge - time='2024-10-01/2025-09-30' + monitoring_location_id="USGS-01646500", + parameter_code="00060", # Discharge + time="2024-10-01/2025-09-30", ) print(f"Retrieved {len(df)} records") @@ -72,9 +73,9 @@ Retrieve streamflow at multiple locations from October 1, 2024 to the present: ```python df, metadata = waterdata.get_daily( - monitoring_location_id=["USGS-13018750","USGS-13013650"], - parameter_code='00060', - time='2024-10-01/..' + monitoring_location_id=["USGS-13018750", "USGS-13013650"], + parameter_code="00060", + time="2024-10-01/..", ) print(f"Retrieved {len(df)} records") @@ -85,8 +86,8 @@ stream sites in Maryland: ```python # Get monitoring location information df, metadata = waterdata.get_monitoring_locations( - state='Maryland', # full name, postal code ('MD'), or FIPS ('24') - site_type_code='ST' # Stream sites + state="Maryland", # full name, postal code ('MD'), or FIPS ('24') + site_type_code="ST", # Stream sites ) print(f"Found {len(df)} stream monitoring locations in Maryland") @@ -98,9 +99,9 @@ windows to avoid timeouts and other issues: ```python # Get continuous data for a single monitoring location and water year df, metadata = waterdata.get_continuous( - monitoring_location_id='USGS-01646500', - parameter_code='00065', # Gage height - time='2024-10-01/2025-09-30' + monitoring_location_id="USGS-01646500", + parameter_code="00065", # Gage height + time="2024-10-01/2025-09-30", ) print(f"Retrieved {len(df)} continuous gage height measurements") ``` @@ -125,10 +126,10 @@ from dataretrieval import waterdata # enough to span many pages, so it profits from a finer split. sites, _ = waterdata.get_monitoring_locations(state="Ohio", site_type_code="ST") -with waterdata.parallel_chunks(32): # fan out into 32 sub-requests +with waterdata.parallel_chunks(32): # fan out into 32 sub-requests df, md = waterdata.get_daily( monitoring_location_id=sites["monitoring_location_id"], - parameter_code="00060", # discharge + parameter_code="00060", # discharge time="2004-01-01/2023-12-31", ) ``` @@ -167,6 +168,7 @@ API — enable debug-level ```python import logging + logging.basicConfig(level=logging.DEBUG) ``` @@ -181,14 +183,14 @@ from dataretrieval import ngwmn # Find the groundwater monitoring sites in a state # (state accepts a full name, a postal code like 'WI', or a FIPS code like '55') -sites, metadata = ngwmn.get_sites(state='Wisconsin') +sites, metadata = ngwmn.get_sites(state="Wisconsin") print(f"Found {len(sites)} NGWMN sites in Wisconsin") # Pull water levels from the first twenty sites over a time window. water_levels, metadata = ngwmn.get_water_level( - monitoring_location_id=sites['monitoring_location_id'][:20], - datetime=['2022-01-01', '2024-01-01'] + monitoring_location_id=sites["monitoring_location_id"][:20], + datetime=["2022-01-01", "2024-01-01"], ) print(f"Retrieved {len(water_levels)} water-level observations") @@ -203,16 +205,15 @@ from dataretrieval import wqp # Find water quality monitoring sites (returns a DataFrame and metadata) sites, metadata = wqp.what_sites( - statecode='US:55', # Wisconsin - siteType='Stream' + statecode="US:55", # Wisconsin + siteType="Stream", ) print(f"Found {len(sites)} stream monitoring sites in Wisconsin") # Get water quality results results, metadata = wqp.get_results( - siteid='USGS-05427718', - characteristicName='Temperature, water' + siteid="USGS-05427718", characteristicName="Temperature, water" ) print(f"Retrieved {len(results)} temperature measurements") @@ -227,18 +228,18 @@ from dataretrieval import nldi # Get watershed basin for a stream reach basin = nldi.get_basin( - feature_source='comid', - feature_id='13293474' # NHD reach identifier + feature_source="comid", + feature_id="13293474", # NHD reach identifier ) print(f"Basin contains {len(basin)} feature(s)") # Find upstream flowlines flowlines = nldi.get_flowlines( - feature_source='comid', - feature_id='13293474', - navigation_mode='UT', # Upstream tributaries - distance=50 # km + feature_source="comid", + feature_id="13293474", + navigation_mode="UT", # Upstream tributaries + distance=50, # km ) print(f"Found {len(flowlines)} upstream tributaries within 50km") @@ -255,17 +256,17 @@ from dataretrieval import wateruse # Monthly public-supply withdrawals for Rhode Island, split into # groundwater and surface-water sources (returns a DataFrame and metadata). df, metadata = wateruse.get_wateruse( - model='wu-public-supply-wd', - variable=['pswdtot', 'pswdgw', 'pswdsw'], - state='RI', # name/postal/FIPS; pass a list to fan out over several areas - start_date='2020-01', - time_resolution='monthly', + model="wu-public-supply-wd", + variable=["pswdtot", "pswdgw", "pswdsw"], + state="RI", # name/postal/FIPS; pass a list to fan out over several areas + start_date="2020-01", + time_resolution="monthly", ) print(f"Retrieved {len(df)} records across {df['huc12_id'].nunique()} watersheds") # Aggregate the HUC12 grid to a statewide monthly total (million gallons/day) -statewide = df.groupby('year_month')['pswdtot_mgd'].sum() +statewide = df.groupby("year_month")["pswdtot_mgd"].sum() print(statewide.head()) ``` diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index d0dae01c..b46c17df 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -74,7 +74,7 @@ import asyncio import functools import os -from collections.abc import Callable, Iterator +from collections.abc import Awaitable, Callable, Iterator from contextlib import contextmanager from contextvars import copy_context from typing import Any, cast @@ -86,17 +86,14 @@ from dataretrieval.utils import HTTPX_DEFAULTS, Ambient, _require_positive_int from . import progress as _progress -from .interruptions import ( - ChunkInterrupted, - _Fetch, - _Finalize, - _passthrough_result, -) -from .planning import ( - ChunkPlan, +from .combining import ( _combine_chunk_frames, _combine_chunk_responses, ) +from .interruptions import ( + ChunkInterrupted, +) +from .planning import ChunkPlan from .retry import ( _NO_RETRY, RetryPolicy, @@ -291,6 +288,30 @@ def parallel_chunks(n: int) -> Iterator[None]: yield +# --------------------------------------------------------------------------- +# Type aliases for the ChunkedCall contract. +# --------------------------------------------------------------------------- + +# The per-sub-request fetcher the decorator wraps and ``ChunkedCall`` drives: +# an ``async def fetch(args) -> (df, response)``. +_Fetch = Callable[[dict[str, Any]], Awaitable[tuple[pd.DataFrame, httpx.Response]]] + +# Caller-supplied transform applied to the combined chunk result, so a +# resumed call returns the same shape as an un-interrupted one rather than +# the chunker's raw ``(frame, httpx.Response)``. This keeps the chunker +# generic: the OGC getters inject their post-processing (type coercion, +# column arrangement, ``BaseMetadata``) through ``_finalize_ogc``. +# The default is identity, so direct ``ChunkedCall`` use is unaffected. +_Finalize = Callable[[pd.DataFrame, httpx.Response], tuple[pd.DataFrame, Any]] + + +def _passthrough_result( + frame: pd.DataFrame, response: httpx.Response +) -> tuple[pd.DataFrame, Any]: + """Default :data:`_Finalize`: return the raw combined pair unchanged.""" + return frame, response + + class ChunkedCall: """ Stateful handle for a chunked call. @@ -417,12 +438,11 @@ def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]: Frames concatenate in sub-args *index* order (``sorted`` keys — deterministic, independent of parallel completion order). The - aggregated response takes its headers from the most-recently- - *completed* sub-request: the ``track`` closure in :meth:`_run` - is the only writer of ``self._chunks`` and ``dict`` preserves - insertion order, so the chunks' natural order is completion - order and the last one carries the freshest - ``x-ratelimit-remaining``. + aggregated response takes its headers from the response with the + lowest reported ``x-ratelimit-remaining`` value. If no response + reports that header, it falls back to the last completed response; + ``self._chunks`` preserves completion order because the ``track`` + closure in :meth:`_run` is its only writer. Returns ------- @@ -521,8 +541,9 @@ def resume(self) -> tuple[pd.DataFrame, Any]: Combined data from every successful sub-request. response The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, most-recently-completed sub-request's headers, - cumulative elapsed time) by default, or whatever + (canonical URL, headers from the response with the lowest reported + remaining quota, and summed response elapsed durations) by default, + or whatever :attr:`finalize` produces (e.g. ``BaseMetadata`` for the OGC getters). @@ -603,8 +624,9 @@ async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: Combined data from every sub-request. response The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, most-recently-completed sub-request's headers, - cumulative elapsed time) by default, or whatever + (canonical URL, headers from the response with the lowest reported + remaining quota, and summed response elapsed durations) by default, + or whatever :attr:`finalize` produces (e.g. ``BaseMetadata`` for OGC getters). Raises diff --git a/dataretrieval/ogc/combining.py b/dataretrieval/ogc/combining.py new file mode 100644 index 00000000..be4366c1 --- /dev/null +++ b/dataretrieval/ogc/combining.py @@ -0,0 +1,206 @@ +"""Result recombination: merge per-chunk frames and responses (no I/O). + +These utilities assemble the output of a chunked/fan-out call from its +individual per-sub-request results. They have no event-loop, retry, or +network state — they're pure data transforms imported by both the +chunked-call execution (:mod:`dataretrieval.ogc.chunking`) and the +per-page pagination (:mod:`dataretrieval.ogc.engine`). + +Separated from :mod:`dataretrieval.ogc.planning` so that module stays +focused on *what* to split, while this module owns *how* to reassemble. +""" + +from __future__ import annotations + +import copy +from datetime import timedelta + +import httpx +import pandas as pd + +# Response header USGS uses to advertise remaining hourly quota. Lives in this +# module so every layer (the combine helpers below, the engine's per-page +# progress reporter) reads it from one place rather than hard-coding the string. +_QUOTA_HEADER = "x-ratelimit-remaining" + + +def _safe_elapsed(response: httpx.Response) -> timedelta: + """ + Read ``response.elapsed``, falling back to ``timedelta(0)`` when + the attribute hasn't been populated. + + httpx only writes ``.elapsed`` when a response is closed through + its normal transport path. ``MockTransport`` (used by + ``pytest-httpx``) and hand-constructed ``httpx.Response`` objects + leave the attribute unset, so accessing it raises ``RuntimeError``. + Combining responses across chunks needs a defined duration, so we + treat the missing attribute as zero elapsed. + """ + try: + elapsed: object = response.elapsed + except RuntimeError: + return timedelta(0) + return elapsed if isinstance(elapsed, timedelta) else timedelta(0) + + +def _set_response_url(response: httpx.Response, url: str | httpx.URL) -> None: + """ + Overwrite the URL surfaced by a response without back-propagating + the change into any aliased original. + + Lightweight test doubles expose ``.url`` as a writable attribute. Real + :class:`httpx.Response` objects resolve it through a bound request, so swap + in a fresh request carrying the new URL; mutating the existing request would + leak through any shallow copy that shares it. + """ + if not isinstance(response, httpx.Response): + # Lightweight test doubles expose ``url`` as a writable attribute. + response.url = url + return + + target = httpx.URL(str(url)) + try: + old = response.request + except RuntimeError: + # No request bound (some hand-built httpx.Response fixtures); + # synthesize a minimal one to hold the URL. + response.request = httpx.Request("GET", target) + return + response.request = httpx.Request(method=old.method, url=target, headers=old.headers) + + +def _lowest_remaining(responses: list[httpx.Response]) -> httpx.Response: + """The response reporting the lowest ``x-ratelimit-remaining``. + + Within a rate-limit window, the counter decreases monotonically, so the + smallest value observed is the most conservative value to surface. Under + concurrent fan-out, the last response *by index* need not be the one the + server processed last. Fall back to the last response when none reports + the header. + """ + best: httpx.Response | None = None + best_remaining: int | None = None + for response in responses: + try: + remaining = int(response.headers[_QUOTA_HEADER]) + except (KeyError, ValueError): + continue + if best_remaining is None or remaining < best_remaining: + best, best_remaining = response, remaining + return best if best is not None else responses[-1] + + +def _merge_response( + base: httpx.Response, + *, + headers_from: httpx.Response, + elapsed: timedelta, + url: str | httpx.URL | None = None, +) -> httpx.Response: + """Fold several responses into one: a shallow copy of ``base`` whose + ``.headers`` are rebuilt as a fresh ``httpx.Headers`` from ``headers_from``, + ``.elapsed`` set to ``elapsed``, and ``.url`` overridden when ``url`` is + given. ``base`` and ``headers_from`` are never mutated, and the fresh + ``httpx.Headers`` means downstream mutations don't back-propagate into any + underlying response — so callers may re-fold idempotently. This is the one + low-level merge behind both pagination + (:func:`~dataretrieval.ogc.engine._paginate`) and the chunked / fan-out + aggregation (:func:`_combine_chunk_responses`).""" + merged = copy.copy(base) + merged.headers = httpx.Headers(headers_from.headers) + merged.elapsed = elapsed + if url is not None: + _set_response_url(merged, url) + return merged + + +def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame: + """Concatenate per-chunk frames and deduplicate IDs across chunks. + + Empty frames are ignored before concatenation so an empty plain + :class:`pandas.DataFrame` cannot downgrade a real ``GeoDataFrame`` and + strip its geometry or CRS. When every frame is empty, the first frame is + returned to preserve its concrete type. + + When multiple non-empty frames are combined, non-null feature IDs are + deduplicated regardless of the plan axis. Filter clauses can match the same + feature, and list inputs can contain repeated values or otherwise select + overlapping records. Rows without an ``id`` are preserved verbatim: pandas + treats null values as duplicates, so deduplicating them would silently lose + data. + """ + non_empty = [frame for frame in frames if not frame.empty] + if not non_empty: + return frames[0] if frames else pd.DataFrame() + if len(non_empty) == 1: + return non_empty[0].copy() + + combined = pd.concat(non_empty, ignore_index=True) + if "id" not in combined.columns: + return combined + + has_id = combined["id"].notna() + if has_id.all(): + return combined.drop_duplicates(subset="id", ignore_index=True) + if has_id.any(): + id_rows = combined[has_id].drop_duplicates(subset="id") + no_id_rows = combined[~has_id] + return pd.concat([id_rows, no_id_rows], ignore_index=True) + return combined + + +def _combine_chunk_responses( + responses: list[httpx.Response], canonical_url: str | None +) -> httpx.Response: + """ + Fold per-sub-request responses into a single aggregated response. + + For a multi-response input, returns a shallow copy of + ``responses[0]`` with ``.headers`` set to those of the response reporting + the lowest ``x-ratelimit-remaining`` value (the most conservative quota + observation; see :func:`_lowest_remaining`), ``.elapsed`` set to the sum of + the per-response elapsed durations, and ``.url`` set to the + canonical original-query URL (when supplied) so ``BaseMetadata`` + reflects the user's full request rather than the first chunk. + + For a single-response input with no canonical-URL override, + ``responses[0]`` is returned unchanged to skip the copy on the + passthrough hot path. + + Parameters + ---------- + responses : list[httpx.Response] + One response per completed sub-request, in caller-provided order. + canonical_url : str or None + URL of the unchunked original request. ``None`` skips the URL + override — used by the passthrough path (the fetcher's + response already carries the original-query URL) and by the + worst-case overflow path (no buildable canonical URL exists). + + Returns + ------- + httpx.Response + A shallow copy of the first response with aggregated + ``headers``, ``elapsed``, and ``url``. The function is + idempotent (the input responses' ``headers`` / ``elapsed`` / + ``url`` are never mutated), so it's safe to call repeatedly + via :attr:`ChunkedCall.partial_response` during error + inspection or resume retries. ``headers`` on the returned + object is a fresh ``httpx.Headers``, so mutations there don't + back-propagate into any chunk's underlying response. + """ + if len(responses) == 1 and canonical_url is None: + return responses[0] + + # Headers come from the response with the lowest reported remaining quota; + # ``_lowest_remaining`` returns the lone response as-is + # for a single-element list). ``_merge_response`` re-sums elapsed onto a + # fresh copy, so repeated calls (e.g. via ``ChunkedCall.partial_response`` + # during resume) stay idempotent. + elapsed = sum((_safe_elapsed(r) for r in responses), start=timedelta()) + return _merge_response( + responses[0], + headers_from=_lowest_remaining(responses), + elapsed=elapsed, + url=canonical_url, + ) diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 668fa7b7..5fda67e5 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -47,9 +47,9 @@ from dataretrieval.ogc import chunking from dataretrieval.ogc import progress as _progress from dataretrieval.ogc.chunking import get_active_client +from dataretrieval.ogc.combining import _QUOTA_HEADER, _merge_response, _safe_elapsed from dataretrieval.ogc.dates import _DATE_RANGE_PARAMS, _format_api_dates from dataretrieval.ogc.errors import _paginated_failure_message, _raise_for_non_200 -from dataretrieval.ogc.planning import _QUOTA_HEADER, _merge_response, _safe_elapsed from dataretrieval.ogc.shaping import GEOPANDAS, _finalize_ogc, _get_resp_data from dataretrieval.utils import ( HTTPX_DEFAULTS, @@ -658,9 +658,9 @@ async def _paginate( response : httpx.Response A shallow copy of the first-page response, with ``.headers`` rebuilt as a fresh ``httpx.Headers`` reflecting the last page and - ``.elapsed`` set to cumulative wall-clock. The canonical URL is - preserved from the first page. The original first-page response - is not mutated. + ``.elapsed`` set to the sum of the per-page response durations. The + canonical URL is preserved from the first page. The original first-page + response is not mutated. Raises ------ diff --git a/dataretrieval/ogc/interruptions.py b/dataretrieval/ogc/interruptions.py index 18398654..8cb5723c 100644 --- a/dataretrieval/ogc/interruptions.py +++ b/dataretrieval/ogc/interruptions.py @@ -11,7 +11,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any, ClassVar import httpx @@ -23,27 +22,6 @@ from dataretrieval.ogc.chunking import ChunkedCall -# ``_Fetch`` is the per-sub-request fetcher the decorator wraps and -# ``ChunkedCall`` drives: an ``async def fetch(args) -> (df, response)``. -_Fetch = Callable[[dict[str, Any]], Awaitable[tuple[pd.DataFrame, httpx.Response]]] - - -# Caller-supplied transform applied to the combined chunk result, so a -# resumed call returns the same shape as an un-interrupted one rather than -# the chunker's raw ``(frame, httpx.Response)``. This keeps the chunker -# generic: the OGC getters inject their post-processing (type coercion, -# column arrangement, ``BaseMetadata``) through ``utils._finalize_ogc``. -# The default is identity, so direct ``ChunkedCall`` use is unaffected. -_Finalize = Callable[[pd.DataFrame, httpx.Response], tuple[pd.DataFrame, Any]] - - -def _passthrough_result( - frame: pd.DataFrame, response: httpx.Response -) -> tuple[pd.DataFrame, Any]: - """Default :data:`_Finalize`: return the raw combined pair unchanged.""" - return frame, response - - class ChunkInterrupted(DataRetrievalError): """ Base class for mid-stream chunk failures whose completed work is @@ -140,10 +118,14 @@ def __init__( self.total_chunks = total_chunks self.call = call self.retry_after = retry_after - # Snapshot partial state at raise time so the exception's view stays - # stable across later ``call.resume()`` advances (the live view is on - # ``call.partial_frame`` / ``.partial_response``). ``.copy()`` guards - # the single-chunk fast path, where the frame may be returned verbatim. + # Snapshot partial state at raise time so the exception stays a stable + # record of the failure moment: ``exc.partial_frame`` / + # ``.partial_response`` do NOT advance on a later ``call.resume()`` + # (that live view is on ``call.partial_frame`` / ``.partial_response``). + # This keeps each interruption in a resume loop a faithful record of + # what it saw, rather than every exception aliasing the shared call's + # advancing state. ``.copy()`` guards the single-chunk fast path, where + # the combined frame may be returned verbatim. if call is None: self.partial_frame: pd.DataFrame = pd.DataFrame() self.partial_response: httpx.Response | None = None @@ -155,9 +137,10 @@ def __getstate__(self) -> dict[str, Any]: # Drop the live ChunkedCall before pickling: its ``.fetch`` is an # undecorated module function pickle can't reference by name, so the # interruption can't cross a process boundary with ``.call`` attached. - # The degraded ``call=None`` form keeps the counts, retry hint, and - # partial frame / response; only ``.resume()`` is lost (cross-process - # resume was never possible anyway). + # The degraded ``call=None`` form keeps the counts, retry hint, and the + # snapshotted partial frame / response — plain instance attributes the + # base ``__getstate__`` already pickles; only ``.resume()`` is lost + # (cross-process resume was never possible anyway). return {**super().__getstate__(), "call": None} diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index e80d0643..15b397a0 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -1,31 +1,32 @@ -"""Pure URL-byte chunk planning and result recombination (no I/O). - -This module holds the side-effect-free half of the chunker: deciding how -to split one over-budget OGC request into URL-fitting sub-requests -(:class:`ChunkPlan` and the axis/byte-accounting helpers) and reassembling -their per-chunk frames and responses (:func:`_combine_chunk_frames`, -:func:`_combine_chunk_responses`). It has no event loop, retry policy, or -network state — those live in :mod:`dataretrieval.ogc.chunking` (execution) -and :mod:`dataretrieval.ogc.retry` (retry policy), which import the plan and -drive it. Keeping the planning/combination logic here -makes it unit-testable without an HTTP client and gives the two concerns -separate reasons to change. +"""Pure URL-byte chunk planning (no I/O). + +This module holds the side-effect-free planning half of the chunker: +deciding how to split one over-budget OGC request into URL-fitting +sub-requests (:class:`ChunkPlan` and the axis/byte-accounting helpers). +It has no event loop, retry policy, or network state — those live in +:mod:`dataretrieval.ogc.chunking` (execution) and +:mod:`dataretrieval.ogc.retry` (retry policy), which import the plan and +drive it. + +Result recombination — reassembling the per-chunk frames and responses +back into one result +(:func:`~dataretrieval.ogc.combining._combine_chunk_frames`, +:func:`~dataretrieval.ogc.combining._combine_chunk_responses`, etc.) — lives in +the sibling :mod:`dataretrieval.ogc.combining` module, which callers import +directly. """ from __future__ import annotations -import copy import itertools import math from collections.abc import Callable, Iterator from contextlib import suppress from dataclasses import dataclass -from datetime import timedelta from typing import Any from urllib.parse import quote_plus import httpx -import pandas as pd from dataretrieval.exceptions import Unchunkable from dataretrieval.ogc.filters import ( @@ -130,52 +131,6 @@ def _safe_request_bytes( return _request_bytes(req) -def _safe_elapsed(response: httpx.Response) -> timedelta: - """ - Read ``response.elapsed``, falling back to ``timedelta(0)`` when - the attribute hasn't been populated. - - httpx only writes ``.elapsed`` when a response is closed through - its normal transport path. ``MockTransport`` (used by - ``pytest-httpx``) and hand-constructed ``httpx.Response`` objects - leave the attribute unset, so accessing it raises ``RuntimeError``. - Combining responses across chunks needs a defined duration, so we - treat the missing attribute as zero elapsed. - """ - try: - return response.elapsed - except RuntimeError: - return timedelta(0) - - -def _set_response_url(response: httpx.Response, url: str | httpx.URL) -> None: - """ - Overwrite the URL surfaced by a response without back-propagating - the change into any aliased original. - - Try the direct assignment first: on lightweight test mocks ``.url`` - is a plain writable attribute. On real ``httpx.Response`` it's - read-only (it resolves through the bound request), so swap in a - fresh :class:`httpx.Request` carrying the new URL — mutating the - existing one would leak through any shallow copy that shares the - same ``.request``. - """ - try: - response.url = url # type: ignore[misc, assignment] - except AttributeError: - target = httpx.URL(str(url)) - try: - old = response.request - except RuntimeError: - # No request bound (some hand-built httpx.Response fixtures); - # synthesize a minimal one to hold the URL. - response.request = httpx.Request("GET", target) - return - response.request = httpx.Request( - method=old.method, url=target, headers=old.headers - ) - - @dataclass(frozen=True) class _Axis: """ @@ -619,171 +574,3 @@ def iter_sub_args(self) -> Iterator[dict[str, Any]]: for axis, chunk in zip(self.axes, combo, strict=False): sub_args[axis.arg_key] = axis.render(chunk) yield sub_args - - -def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame: - """ - Concatenate per-chunk frames, dropping empties and deduping by ``id``. - - Parameters - ---------- - frames : list[pandas.DataFrame] - One frame per completed sub-request. - - Returns - ------- - pandas.DataFrame - The concatenated, deduplicated result. Empty when every input - frame is empty. - - Notes - ----- - An empty chunk can be a plain ``pd.DataFrame()`` (no geopandas); - concatenating it with real ``GeoDataFrame``s downgrades the result - to plain ``DataFrame`` and strips geometry/CRS, so empties are - dropped first. Dedup on the pre-rename feature ``id`` keeps - overlapping user OR-clauses from producing duplicate rows across - chunks. - - Dedup is restricted to rows whose ``id`` is non-null. ``pandas`` - treats NaN==NaN as a duplicate for ``drop_duplicates``, so a - blanket call would collapse every id-less row into a single one — - silent data loss if any chunk emits features without an - ``id`` field. - """ - non_empty = [f for f in frames if not f.empty] - if not non_empty: - # Preserve the frame type (GeoDataFrame vs DataFrame) of the - # input even when every chunk is empty — ``_get_resp_data`` - # returns ``gpd.GeoDataFrame()`` on empty geopd responses, and - # returning a plain ``pd.DataFrame()`` here would downgrade - # the type in a downstream ``pd.concat([result, geo_page])`` to - # a plain DataFrame and strip geometry/CRS. - return frames[0] if frames else pd.DataFrame() - if len(non_empty) == 1: - # Single-completed-chunk fast path. Return a copy so callers - # who treat ``ChunkedCall.partial_frame`` as a fresh result - # (the property docstring says "live; recomputed per access") - # don't accidentally mutate ``_chunks[0][0]`` in place. - return non_empty[0].copy() - combined = pd.concat(non_empty, ignore_index=True) - if "id" in combined.columns: - has_id = combined["id"].notna() - if has_id.all(): - combined = combined.drop_duplicates(subset="id", ignore_index=True) - elif has_id.any(): - # Mixed: dedupe only the id-bearing rows; preserve id-less - # rows verbatim (their order relative to id-bearing rows - # may shift, which is acceptable — dedup can't be id-keyed - # for rows without an id). - id_rows = combined[has_id].drop_duplicates(subset="id") - no_id_rows = combined[~has_id] - combined = pd.concat([id_rows, no_id_rows], ignore_index=True) - return combined - - -# Response header USGS uses to advertise remaining hourly quota. Lives in this -# base module so every layer (planning's ``_lowest_remaining``, the engine's -# per-page progress) reads it from one place rather than hard-coding the string. -_QUOTA_HEADER = "x-ratelimit-remaining" - - -def _lowest_remaining(responses: list[httpx.Response]) -> httpx.Response: - """The response reporting the lowest ``x-ratelimit-remaining``. - - The rate-limit counter decreases monotonically within a window, so the - smallest value any sub-request saw is the most-current "quota left after - this call" — the right thing to surface. Under concurrent fan-out the - last response *by index* need not be the one the server processed last, so - pick the minimum (falling back to the last response if none report it). - """ - best: httpx.Response | None = None - best_remaining: int | None = None - for response in responses: - try: - remaining = int(response.headers[_QUOTA_HEADER]) - except (KeyError, ValueError): - continue - if best_remaining is None or remaining < best_remaining: - best, best_remaining = response, remaining - return best if best is not None else responses[-1] - - -def _merge_response( - base: httpx.Response, - *, - headers_from: httpx.Response, - elapsed: timedelta, - url: str | httpx.URL | None = None, -) -> httpx.Response: - """Fold several responses into one: a shallow copy of ``base`` whose - ``.headers`` are rebuilt as a fresh ``httpx.Headers`` from ``headers_from``, - ``.elapsed`` set to ``elapsed``, and ``.url`` overridden when ``url`` is - given. ``base`` and ``headers_from`` are never mutated, and the fresh - ``httpx.Headers`` means downstream mutations don't back-propagate into any - underlying response — so callers may re-fold idempotently. This is the one - low-level merge behind both pagination (:func:`_paginate`) and the chunked / - fan-out aggregation (:func:`_combine_chunk_responses`).""" - merged = copy.copy(base) - merged.headers = httpx.Headers(headers_from.headers) - merged.elapsed = elapsed - if url is not None: - _set_response_url(merged, url) - return merged - - -def _combine_chunk_responses( - responses: list[httpx.Response], canonical_url: str | None -) -> httpx.Response: - """ - Fold per-sub-request responses into a single aggregated response. - - For a multi-response input, returns a shallow copy of - ``responses[0]`` with ``.headers`` set to those of the most-depleted - response (lowest ``x-ratelimit-remaining`` — the quota actually left - after the fan-out; see :func:`_lowest_remaining`), ``.elapsed`` set - to total wall-clock across every response, and ``.url`` set to the - canonical original-query URL (when supplied) so ``BaseMetadata`` - reflects the user's full request rather than the first chunk. - - For a single-response input with no canonical-URL override, - ``responses[0]`` is returned unchanged to skip the copy on the - passthrough hot path. - - Parameters - ---------- - responses : list[httpx.Response] - One response per completed sub-request, in execution order. - canonical_url : str or None - URL of the unchunked original request. ``None`` skips the URL - override — used by the passthrough path (the fetcher's - response already carries the original-query URL) and by the - worst-case overflow path (no buildable canonical URL exists). - - Returns - ------- - httpx.Response - A shallow copy of the first response with aggregated - ``headers``, ``elapsed``, and ``url``. The function is - idempotent (the input responses' ``headers`` / ``elapsed`` / - ``url`` are never mutated), so it's safe to call repeatedly - via :attr:`ChunkedCall.partial_response` during error - inspection or resume retries. ``headers`` on the returned - object is a fresh ``httpx.Headers``, so mutations there don't - back-propagate into any chunk's underlying response. - """ - if len(responses) == 1 and canonical_url is None: - return responses[0] - - # Headers come from the most-depleted response (lowest quota left after a - # concurrent fan-out; ``_lowest_remaining`` returns the lone response as-is - # for a single-element list). ``_merge_response`` re-sums elapsed onto a - # fresh copy, so repeated calls (e.g. via ``ChunkedCall.partial_response`` - # during resume) stay idempotent. - elapsed = sum((_safe_elapsed(r) for r in responses), start=timedelta()) - return _merge_response( - responses[0], - headers_from=_lowest_remaining(responses), - elapsed=elapsed, - url=canonical_url, - ) diff --git a/dataretrieval/ogc/retry.py b/dataretrieval/ogc/retry.py index b81ce60e..37ea040a 100644 --- a/dataretrieval/ogc/retry.py +++ b/dataretrieval/ogc/retry.py @@ -19,7 +19,7 @@ import httpx import pandas as pd -from dataretrieval.exceptions import RateLimited, ServiceUnavailable, TransientError +from dataretrieval.exceptions import RateLimited, TransientError from dataretrieval.ogc import progress as _progress from dataretrieval.ogc.interruptions import ( ChunkInterrupted, @@ -190,6 +190,28 @@ def backoff(self, attempt: int, retry_after: float | None) -> float: _NO_RETRY = RetryPolicy(max_retries=0) +def _classify_transient( + exc: BaseException, +) -> tuple[type[ChunkInterrupted], float | None] | None: + """Classify one exception as a transient, resumable failure. + + This function owns the shared exception taxonomy; it deliberately does not + walk ``__cause__``. :func:`_classify_chunk_error` walks wrapped pagination + failures, while :func:`_retryable` applies the narrower automatic-retry + policy to this classification. + """ + if isinstance(exc, RateLimited): + return QuotaExhausted, exc.retry_after + if isinstance(exc, TransientError): + # Every typed transient other than a rate-limit error is a service + # interruption. This fallback keeps future TransientError subclasses + # resumable after their inline retries are exhausted. + return ServiceInterrupted, exc.retry_after + if isinstance(exc, (httpx.HTTPError, httpx.InvalidURL)): + return ServiceInterrupted, None + return None + + def _classify_chunk_error( exc: BaseException, ) -> tuple[type[ChunkInterrupted], float | None] | None: @@ -231,41 +253,28 @@ def _classify_chunk_error( """ cur: BaseException | None = exc while cur is not None: - if isinstance(cur, RateLimited): - return QuotaExhausted, cur.retry_after - if isinstance(cur, ServiceUnavailable): - return ServiceInterrupted, cur.retry_after - if isinstance(cur, (httpx.HTTPError, httpx.InvalidURL)): - return ServiceInterrupted, None + result = _classify_transient(cur) + if result is not None: + return result cur = cur.__cause__ return None def _retryable(exc: BaseException) -> tuple[bool, float | None]: - """ - Decide whether ``exc`` is a transient worth an automatic retry. - - Only the *top-level* exception is inspected — unlike - :func:`_classify_chunk_error`, which walks the ``__cause__`` chain. - The distinction matters because ``_paginate`` raises an - initial-request transient (429 / 5xx / :class:`httpx.TransportError`) - *raw*, but wraps a mid-pagination failure as a base ``DataRetrievalError``. - So a raw transient means a sub-request that made no progress and is cheap to - re-issue, whereas a mid-pagination failure is left to escalate to a - resumable :class:`ChunkInterrupted` rather than re-walked from page 1 - (which would re-spend the quota just exhausted). ``httpx.InvalidURL`` - is never retried — a too-long cursor won't fix on a retry. + """Decide whether a top-level transient is worth an automatic retry. - Returns - ------- - tuple[bool, float or None] - ``(retryable, retry_after)`` — the server ``Retry-After`` hint - (seconds) when the transient carried one, else ``None``. + Wrapped mid-pagination failures are not retried from page one; they instead + escalate to a resumable :class:`ChunkInterrupted`. ``httpx.InvalidURL`` and + non-transport ``httpx.HTTPError`` instances are classified as resumable but + excluded from automatic retry by policy. """ - if isinstance(exc, TransientError): - return True, exc.retry_after - if isinstance(exc, httpx.TransportError): - return True, None + classification = _classify_transient(exc) + if classification is None: + return False, None + + _, retry_after = classification + if isinstance(exc, (TransientError, httpx.TransportError)): + return True, retry_after return False, None diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index 4df242f2..5870c63b 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -101,17 +101,13 @@ def _get_resp_data( Notes ----- - The non-geopandas branch builds the frame directly from each - feature's ``properties`` dict, plus the top-level ``id`` and - ``geometry.coordinates`` columns — the ``id`` column is always - added (so the downstream rename to the service-specific output id - works even on an all-None id), while the ``geometry`` column is - added only when at least one feature carries geometry. This skips - the GeoJSON envelope entirely, so - newly-added Feature-level fields (e.g. ``geometry.type`` after - USGS migrated to full GeoJSON geometry objects) can't leak into - the result frame; no reactive drop-list needs maintenance every - time the upstream schema grows. + The non-geopandas branch normalizes each feature's ``properties`` object, + flattening nested dictionaries with an underscore separator, then adds the + top-level ``id`` and a ``geometry`` column containing the coordinates. The + ``id`` column is always added so the downstream service-specific rename + works even when all IDs are missing; ``geometry`` is added only when + coordinates are present. Feature-level envelope fields are deliberately + excluded. """ if body is None: body = resp.json() @@ -128,13 +124,11 @@ def _get_resp_data( return _empty_feature_frame(geopd) if not geopd: - df = pd.json_normalize([f.get("properties") or {} for f in features], sep="_") - # Always materialize the ``id`` column (may be all-None) so - # ``_arrange_cols``'s ``df.rename(columns={"id": output_id})`` - # produces the documented service-specific output_id column - # (daily_id, channel_measurements_id, …) even if the upstream - # response carried no feature-level id. - df["id"] = [f.get("id") for f in features] + properties = [feature.get("properties") or {} for feature in features] + df = pd.json_normalize(properties, sep="_") + # Always materialize the feature-level ID (possibly all-None) so + # ``_arrange_cols`` can perform the documented service-specific rename. + df["id"] = [feature.get("id") for feature in features] _attach_coordinates(df, features) return df @@ -369,7 +363,7 @@ def _finalize_ogc( as :class:`~dataretrieval.utils.BaseMetadata`. Injected into the chunker as its ``finalize`` hook (see - :data:`~dataretrieval.ogc.interruptions._Finalize`) so the + :data:`~dataretrieval.ogc.chunking._Finalize`) so the un-interrupted return *and* a resumed ``ChunkInterrupted.call.resume()`` produce the same post-processed ``(DataFrame, BaseMetadata)`` shape, not the chunker's raw frame and bare ``httpx.Response``. diff --git a/dataretrieval/wateruse.py b/dataretrieval/wateruse.py index 635a97fb..5dd5cb9c 100644 --- a/dataretrieval/wateruse.py +++ b/dataretrieval/wateruse.py @@ -53,8 +53,8 @@ from dataretrieval.codes.states import to_state from dataretrieval.exceptions import DataRetrievalError +from dataretrieval.ogc.combining import _combine_chunk_frames, _combine_chunk_responses from dataretrieval.ogc.engine import _paginate, _run_sync -from dataretrieval.ogc.planning import _combine_chunk_frames, _combine_chunk_responses from dataretrieval.utils import ( HTTPX_DEFAULTS, BaseMetadata, @@ -365,8 +365,9 @@ async def _one(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: results = await asyncio.gather(*(_one(req) for req in requests)) # Reuse the engine's combine helpers: drop empty frames and concat, and fold - # the per-location responses into one (lowest-remaining rate-limit headers + - # cumulative elapsed), keeping the first request's URL as the query identity. + # the per-location responses into one (headers from the response with the + # lowest reported remaining quota plus summed response durations), keeping + # the first request's URL as the query identity. frames = [frame for frame, _ in results] responses = [resp for _, resp in results] return _combine_chunk_frames(frames), _combine_chunk_responses( @@ -398,7 +399,7 @@ def _next_page_url(response: httpx.Response) -> str | None: url = response.links.get("next", {}).get("url") if not url: return None - return url.replace("https://water.usgs.gov", "https://api.water.usgs.gov", 1) + return str(url).replace("https://water.usgs.gov", "https://api.water.usgs.gov", 1) def _nwdc_error_detail(response: httpx.Response) -> str | None: diff --git a/pyproject.toml b/pyproject.toml index e261707b..3e6bfbe7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,7 @@ test = [ "pytest-rerunfailures", "coverage", "pytest-httpx", - "ruff", + "ruff==0.16.1", "dataretrieval[type-check]", # mypy, pinned once in the type-check extra ] doc = [ diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index 600f05d8..4a86bf17 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -35,6 +35,7 @@ DataRetrievalError, RateLimited, ServiceUnavailable, + TransientError, Unchunkable, ) from dataretrieval.ogc import chunking as _chunking @@ -47,6 +48,11 @@ multi_value_chunked, parallel_chunks, ) +from dataretrieval.ogc.combining import ( + _QUOTA_HEADER, + _combine_chunk_frames, + _combine_chunk_responses, +) from dataretrieval.ogc.interruptions import ( ChunkInterrupted, QuotaExhausted, @@ -56,10 +62,7 @@ _LIST_SEP, _NEVER_CHUNK, _OR_SEP, - _QUOTA_HEADER, ChunkPlan, - _combine_chunk_frames, - _combine_chunk_responses, _extract_axes, _request_bytes, _safe_request_bytes, @@ -1124,15 +1127,29 @@ def test_combine_chunk_frames_does_not_collapse_none_ids(): def test_combine_chunk_frames_still_dedupes_overlapping_ids(): - """The original dedup contract — overlapping OR-clause partitions - that produce duplicate-id rows across chunks must still collapse - to one row — has to keep working when ids ARE present.""" + """Duplicate feature IDs collapse regardless of why chunks overlap.""" df_a = pd.DataFrame({"id": ["x", "y"], "val": [1, 2]}) df_b = pd.DataFrame({"id": ["y", "z"], "val": [2, 3]}) combined = _combine_chunk_frames([df_a, df_b]) assert sorted(combined["id"].tolist()) == ["x", "y", "z"] +def test_list_axis_chunks_dedupe_repeated_feature_ids(): + """Repeated list values can select the same feature in separate chunks.""" + + @multi_value_chunked(build_request=_fake_build, url_limit=8000) + async def fetch(args): + return ( + pd.DataFrame({"id": ["feature-1"], "site": [args["sites"][0]]}), + _quota_response(500), + ) + + with parallel_chunks(2): + frame, _ = fetch({"sites": ["A", "A"]}) + + assert frame.to_dict(orient="records") == [{"id": "feature-1", "site": "A"}] + + def test_retry_after_surfaces_on_quota_exhausted(): """If the 429 response includes a ``Retry-After`` header, that delay must travel from the typed transport exception @@ -1997,6 +2014,25 @@ async def fetch(args): assert sorted(df["sites"]) == sorted(sites) # all recovered despite the 429 +def test_future_transient_subclass_remains_resumable(monkeypatch): + """The shared taxonomy handles new typed transients by default.""" + + class MaintenanceWindow(TransientError): + _DEFAULT_STATUS = 503 + + monkeypatch.setenv("API_USGS_RETRIES", "0") + + async def fetch(args): + raise MaintenanceWindow("planned maintenance", retry_after=30.0) + + decorated = multi_value_chunked(build_request=_fake_build, url_limit=240)(fetch) + with pytest.raises(ServiceInterrupted) as excinfo: + decorated({"sites": ["S1" * 10, "S2" * 10]}) + + assert excinfo.value.retry_after == 30.0 + assert isinstance(excinfo.value.__cause__, MaintenanceWindow) + + def test_chunker_exhausted_retries_still_resumable(monkeypatch): """When retries are exhausted the failure still surfaces as a resumable ChunkInterrupted — retries don't swallow the escape hatch.""" diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index 95aabe30..431cf7f4 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -676,6 +676,25 @@ def test_get_resp_data_always_materializes_id_column(): assert df["id"].isna().all() +def test_get_resp_data_flattens_nested_properties(): + """Nested GeoJSON properties keep underscore-separated column names.""" + resp = mock.MagicMock() + resp.json.return_value = { + "features": [ + { + "id": "feature-1", + "properties": {"station": {"code": "A"}, "value": 1}, + } + ] + } + + df = _get_resp_data(resp, geopd=False) + + assert df.to_dict(orient="records") == [ + {"value": 1, "station_code": "A", "id": "feature-1"} + ] + + # --- _arrange_cols ---------------------------------------------------------- diff --git a/tests/wateruse_test.py b/tests/wateruse_test.py index 031aad74..00a843c8 100644 --- a/tests/wateruse_test.py +++ b/tests/wateruse_test.py @@ -306,7 +306,7 @@ def test_fan_out_surfaces_final_rate_limit_header(httpx_mock): assert md.header["x-ratelimit-remaining"] == "850" -# (response aggregation now reuses ogc.planning._combine_chunk_responses; the +# (response aggregation now reuses ogc.combining._combine_chunk_responses; the # integration test above pins the rate-limit-header behavior end-to-end.)