refactor(waterdata): /simplify pass — typed RateLimited exception, drop sentinels

Pulls cleanup items out of /simplify's review:

- **Replace string-sniffing 429 detection with a typed exception.**
  ``_raise_for_non_200`` now raises ``utils.RateLimited`` (a
  ``RuntimeError`` subclass) for HTTP 429 specifically, plain
  ``RuntimeError`` for everything else. The chunker's ``_is_429``
  walks ``__cause__`` looking for ``isinstance(_, RateLimited)``
  instead of ``str(cur).startswith("429")``; ``_paginated_failure_message``
  similarly switches to ``isinstance``. Any future reformatting of
  the 429 message can no longer silently break quota handling.

- **Drop the ``_QUOTA_UNKNOWN = -1`` sentinel.** ``_read_remaining``
  now returns ``int | None``; the wrapper branches on
  ``remaining is not None`` instead of comparing against a magic
  number.

- **Collapse ``_finalize_response`` into ``_combine_chunk_responses``.**
  The two were a 3-line wrapper around a 6-line helper; merging
  them removes one call and the indirection. The combiner now takes
  ``canonical_url`` and sets ``.url`` directly.

- **Simplify ``_FetchOnce``** from a constrained ``TypeVar`` to a
  plain ``Callable`` alias. The TypeVar required ``# type: ignore``
  at the return site anyway and bought no callsite type safety.

- Update the ``waterdata_test.py`` flaky-rerun regex to match the
  new ``RateLimited:`` prefix as well as ``RuntimeError:``.

Items considered and skipped:

- Planning-phase efficiency findings (redundant ``build_request``
  probes, ``_filter_candidates`` joining discarded groups, dict-copy
  cost in ``_iter_sub_args``) — all cold-path next to the actual HTTP
  round-trips that follow. Premature optimization.
- Unifying ``_combine_chunk_responses`` with
  ``_finalize_paginated_response`` — they take different inputs
  (one accumulates wall-clock externally, the other builds it from
  the response list); unification would be churn.
- Test docstring trim — separate pass.

All 142 chunker / filter / utils unit tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: thodson-usgs

dataretrieval/waterdata/chunking.py

@@ -34,7 +34,7 @@
 import itertools
 import math
 from collections.abc import Callable, Iterator
-from typing import Any, TypeVar
+from typing import Any
 from urllib.parse import quote_plus
 
 import pandas as pd
@@ -87,12 +87,6 @@
 # Response header USGS uses to advertise remaining hourly quota.
 _QUOTA_HEADER = "x-ratelimit-remaining"
 
-# Sentinel returned by ``_read_remaining`` when the response has no
-# parseable header. The wrapper treats this as "no quota signal" and
-# skips the post-first-chunk plan-vs-quota check, so the value just
-# needs to be distinct from any plausible real remaining count.
-_QUOTA_UNKNOWN = -1
-
 # Separators the two chunking dimensions use to compose their atoms
 # into URL-encoded blobs. List dims comma-join values
 # (``site=USGS-A,USGS-B``); filter dims OR-join clauses
@@ -101,11 +95,7 @@
 _LIST_SEP = ","
 _OR_SEP = " OR "
 
-
-_FetchOnce = TypeVar(
-    "_FetchOnce",
-    bound=Callable[[dict[str, Any]], tuple[pd.DataFrame, requests.Response]],
-)
+_FetchOnce = Callable[[dict[str, Any]], tuple[pd.DataFrame, requests.Response]]
 
 
 class RequestTooLarge(ValueError):
@@ -437,27 +427,34 @@ def _plan_joint(
     return best[1], best[2]
 
 
-def _read_remaining(response: requests.Response) -> int:
-    """Parse ``x-ratelimit-remaining`` from a response. Missing or
-    malformed header → ``_QUOTA_UNKNOWN`` (treated as "no signal"; the
-    post-first-chunk quota check skips when this is returned)."""
+def _read_remaining(response: requests.Response) -> int | None:
+    """Parse ``x-ratelimit-remaining`` from a response. Returns ``None``
+    when the header is missing or unparseable; the wrapper treats that
+    as "no quota signal" and skips the post-first-chunk plan check."""
     raw = response.headers.get(_QUOTA_HEADER)
     if raw is None:
-        return _QUOTA_UNKNOWN
+        return None
     try:
         return int(raw)
     except (TypeError, ValueError):
-        return _QUOTA_UNKNOWN
+        return None
 
 
 def _is_429(exc: BaseException) -> bool:
-    """True iff ``exc`` or any link in its ``__cause__`` chain is a
-    ``RuntimeError`` whose message starts with ``429`` — the shape
-    ``_raise_for_non_200`` produces and ``_walk_pages`` re-wraps for
-    mid-pagination failures."""
+    """True iff ``exc`` or anywhere along its ``__cause__`` chain is a
+    ``utils.RateLimited``. ``_walk_pages`` re-wraps mid-pagination
+    failures as ``RuntimeError`` with the original ``RateLimited``
+    linked as ``__cause__``, so the chunker has to follow the chain
+    rather than just ``isinstance(exc, RateLimited)``.
+
+    Lazy import: ``utils`` imports this module to decorate
+    ``_fetch_once``, so a top-level import would be circular.
+    """
+    from .utils import RateLimited
+
     cur: BaseException | None = exc
     while cur is not None:
-        if isinstance(cur, RuntimeError) and str(cur).startswith("429"):
+        if isinstance(cur, RateLimited):
             return True
         cur = cur.__cause__
     return False
@@ -499,38 +496,22 @@ def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame:
 
 
 def _combine_chunk_responses(
-    responses: list[requests.Response],
+    responses: list[requests.Response], canonical_url: str
 ) -> requests.Response:
-    """Return one response with the last chunk's headers (for current
-    rate-limit state) and summed ``elapsed`` (for total wall-clock).
-
-    The returned response's ``.url`` is the *first chunk's* URL, which
-    only reflects the first slice of the user's query. ``_finalize_response``
-    overwrites ``.url`` with the canonical original-query URL so
-    ``BaseMetadata`` reflects the user's request, not the first sub-chunk.
-
-    Mutates the first response in place: ``.headers`` is replaced with
-    the last response's headers and ``.elapsed`` is accumulated across
-    all chunks. Downstream reads ``.url``, ``.headers``, and
-    ``.elapsed`` (via ``BaseMetadata``).
-    """
+    """Fold per-sub-request responses into one. The first response is
+    mutated in place: ``.headers`` becomes the last response's (so
+    ``x-ratelimit-remaining`` reflects current state), ``.elapsed``
+    accumulates total wall-clock, and ``.url`` is set to the canonical
+    original-query URL so ``BaseMetadata`` reflects the user's full
+    request rather than the first sub-chunk."""
     head = responses[0]
     if len(responses) > 1:
         head.headers = responses[-1].headers
         head.elapsed = sum((r.elapsed for r in responses[1:]), start=head.elapsed)
+    head.url = canonical_url
     return head
 
 
-def _finalize_response(
-    responses: list[requests.Response], canonical_url: str
-) -> requests.Response:
-    """Aggregate per-sub-request responses and restore the canonical
-    URL representing the user's full original query."""
-    combined = _combine_chunk_responses(responses)
-    combined.url = canonical_url
-    return combined
-
-
 def multi_value_chunked(
     *,
     build_request: Callable[..., Any],
@@ -585,7 +566,7 @@ def wrapper(
                     raise QuotaExhausted(
                         partial_frame=_combine_chunk_frames(frames),
                         partial_response=(
-                            _finalize_response(responses, canonical_url)
+                            _combine_chunk_responses(responses, canonical_url)
                             if responses
                             else None
                         ),
@@ -596,7 +577,7 @@ def wrapper(
                 responses.append(response)
                 if i == 0 and total > 1:
                     remaining = _read_remaining(response)
-                    if remaining != _QUOTA_UNKNOWN and remaining < total - 1:
+                    if remaining is not None and remaining < total - 1:
                         raise RequestExceedsQuota(
                             planned_chunks=total,
                             available=remaining + 1,
@@ -605,9 +586,9 @@ def wrapper(
 
             return (
                 _combine_chunk_frames(frames),
-                _finalize_response(responses, canonical_url),
+                _combine_chunk_responses(responses, canonical_url),
             )
 
-        return wrapper  # type: ignore[return-value]
+        return wrapper
 
     return decorator

dataretrieval/waterdata/utils.py

@@ -410,16 +410,30 @@ def _error_body(resp: requests.Response):
     )
 
 
+class RateLimited(RuntimeError):
+    """A USGS Water Data API request was rejected with HTTP 429. Exposed
+    as a typed exception so callers (notably the multi-value chunker)
+    can detect rate-limit failures via ``isinstance`` instead of
+    string-matching error messages."""
+
+
 def _raise_for_non_200(resp: requests.Response) -> None:
-    """Raise ``RuntimeError(_error_body(resp))`` if ``resp`` is not 200.
+    """Raise on a non-200 response. ``RateLimited`` for 429 (so the
+    chunker can branch on it without parsing the message); plain
+    ``RuntimeError`` for every other failure.
 
     Routes through ``_error_body`` (USGS-API-aware: handles 429/403
     specially, extracts ``code``/``description`` from JSON error bodies)
     rather than ``Response.raise_for_status``, which raises
     ``HTTPError`` with a generic message.
     """
-    if resp.status_code != 200:
-        raise RuntimeError(_error_body(resp))
+    status = resp.status_code
+    if status == 200:
+        return
+    body = _error_body(resp)
+    if status == 429:
+        raise RateLimited(body)
+    raise RuntimeError(body)
 
 
 def _paginated_failure_message(pages_collected: int, cause: BaseException) -> str:
@@ -436,7 +450,7 @@ def _paginated_failure_message(pages_collected: int, cause: BaseException) -> st
     # message is always informative.
     if not cause_str.strip():
         cause_str = type(cause).__name__
-    if cause_str.startswith("429"):
+    if isinstance(cause, RateLimited):
         action = "wait for the rate-limit window to reset and retry"
     else:
         action = "retry the request (possibly after a short backoff)"

tests/waterdata_chunking_test.py

@@ -42,7 +42,7 @@
     _read_remaining,
     multi_value_chunked,
 )
-from dataretrieval.waterdata.utils import _construct_api_requests
+from dataretrieval.waterdata.utils import RateLimited, _construct_api_requests
 
 
 class _FakeReq:
@@ -310,21 +310,16 @@ def test_read_remaining_parses_header():
     assert _read_remaining(_quota_response(42)) == 42
 
 
-def test_read_remaining_treats_missing_header_as_unknown():
-    """Servers that don't echo a rate-limit header must not synthesize
-    a remaining count — the wrapper checks for the unknown sentinel
-    explicitly and skips the post-first-chunk plan-vs-quota check."""
-    from dataretrieval.waterdata.chunking import _QUOTA_UNKNOWN
+def test_read_remaining_returns_none_when_header_missing():
+    """No rate-limit header → ``None`` so the wrapper can branch on
+    ``is None`` instead of comparing against a magic sentinel."""
+    assert _read_remaining(_quota_response(None)) is None
 
-    assert _read_remaining(_quota_response(None)) == _QUOTA_UNKNOWN
 
-
-def test_read_remaining_treats_malformed_header_as_unknown():
-    """Defensive: non-integer header value → unknown sentinel, so the
-    quota check is skipped rather than tripping on a parse failure."""
-    from dataretrieval.waterdata.chunking import _QUOTA_UNKNOWN
-
-    assert _read_remaining(_quota_response("not-a-number")) == _QUOTA_UNKNOWN
+def test_read_remaining_returns_none_on_malformed_header():
+    """Non-integer header value → ``None`` so a parse failure doesn't
+    trip the quota check."""
+    assert _read_remaining(_quota_response("not-a-number")) is None
 
 
 def test_request_exceeds_quota_after_first_chunk():
@@ -399,10 +394,11 @@ def fetch(args):
         i = state["i"]
         state["i"] += 1
         if i == 2:
-            inner = RuntimeError("429: Too many requests made.")
+            # Match _walk_pages's wrapping: a generic mid-pagination
+            # RuntimeError with the typed RateLimited as __cause__.
             try:
-                raise inner
-            except RuntimeError as cause:
+                raise RateLimited("429: Too many requests made.")
+            except RateLimited as cause:
                 raise RuntimeError(
                     "Paginated request failed after collecting 0 page(s): "
                     "429: Too many requests made."
@@ -430,7 +426,7 @@ def test_quota_exhausted_on_first_chunk_429_has_no_partial_response():
     any data arrived" from "abort after partial collection"."""
 
     def fetch(args):
-        raise RuntimeError("429: Too many requests made.")
+        raise RateLimited("429: Too many requests made.")
 
     decorated = multi_value_chunked(build_request=_fake_build, url_limit=240)(fetch)
     with pytest.raises(QuotaExhausted) as excinfo:

tests/waterdata_test.py

@@ -49,7 +49,7 @@
     reruns=2,
     reruns_delay=5,
     only_rerun=[
-        r"RuntimeError:\s*(?:429|5\d\d):",  # _raise_for_non_200 output
+        r"(?:RateLimited|RuntimeError):\s*(?:429|5\d\d):",  # _raise_for_non_200 output
         r"ConnectionError",
         r"ReadTimeout|ConnectTimeout|Timeout",
     ],