refactor(waterdata): simplify chunking for global coherence

Quality cleanup of the chunking machinery after its several iterations — no
behavior change (live bit-for-bit chunked-vs-unchunked still identical for the
GET getters; 165 affected tests + README/time-conventions doctests pass).

- Delete the vestigial `ChunkPlan.execute` pass-through (one caller, no tests)
  and inline `ChunkedCall(plan, fetch, postprocess).resume()` into the
  decorator — removes a hop and a ~35-line docstring that duplicated
  ChunkedCall's contract.
- `ChunkedCall.resume()` reads its own `partial_frame` / `partial_response`
  accessors instead of re-deriving the combine inline — one definition of
  "the combined result".
- Add a `_ChunkedFetch` alias so the decorator's return type is honest
  (`multi_value_chunked` no longer claims to return a bare `_FetchOnce`; the
  wrapper takes `*, postprocess` and returns `(df, Any)`).
- Collapse the finalizer rationale to one home (`_PostProcess`); trim the
  repeated copies in `ChunkedCall.__init__` and `get_ogc_data._postprocess`.
- Fix stale/garbled comments: the `copy.copy` snapshot rationale (false on the
  single-response path), the `_combine_chunk_frames` all-empty comment (it
  referenced a removed `pd.concat([result, geo_page])` and had broken
  grammar), and the dropped `ChunkPlan.execute` docstring reference.

Net -45 lines. ruff clean.

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

Author: thodson-usgs

dataretrieval/waterdata/chunking.py

@@ -151,6 +151,11 @@ def get_active_client() -> httpx.Client | None:
 # default) returns the raw combined ``(frame, response)``.
 _PostProcess = Callable[[pd.DataFrame, httpx.Response], tuple[pd.DataFrame, Any]]
 
+# The callable ``multi_value_chunked`` returns: like ``_FetchOnce`` but with an
+# optional keyword-only ``postprocess`` finalizer, returning either the raw
+# combined ``(frame, response)`` or the finalizer's output.
+_ChunkedFetch = Callable[..., tuple[pd.DataFrame, Any]]
+
 
 class _RetryableTransportError(RuntimeError):
     """
@@ -320,10 +325,11 @@ def __init__(
         # empty-frame schema fetch, or BaseMetadata reading an unset
         # ``elapsed``) could mask this interruption and destroy the
         # resumable ``.call`` handle.
-        # ``partial_frame`` gets a defensive ``.copy()`` because
-        # ``_combine_chunk_frames`` may return a chunk frame verbatim
-        # in the single-completed-chunk fast path; ``partial_response``
-        # already comes via ``copy.copy`` from ``_combine_chunk_responses``.
+        # ``partial_frame`` gets a defensive ``.copy()`` so the snapshot is
+        # independent of the live combine (``_combine_chunk_frames`` can
+        # return an underlying chunk frame verbatim). The aggregated
+        # ``partial_response`` is only ever read, never mutated, so it
+        # needs no copy.
         if call is None:
             self.partial_frame: pd.DataFrame = pd.DataFrame()
             self.partial_response: httpx.Response | None = None
@@ -784,46 +790,6 @@ def iter_sub_args(self) -> Iterator[dict[str, Any]]:
                 sub_args[axis.arg_key] = axis.render(chunk)
             yield sub_args
 
-    def execute(
-        self,
-        fetch_once: _FetchOnce,
-        postprocess: _PostProcess | None = None,
-    ) -> tuple[pd.DataFrame, Any]:
-        """
-        Run the plan and return the combined result.
-
-        Thin wrapper around
-        ``ChunkedCall(self, fetch_once, postprocess).resume()``; see
-        :class:`ChunkedCall` for the per-sub-request semantics.
-
-        Parameters
-        ----------
-        fetch_once : Callable
-            Function that issues a single sub-request, given the
-            substituted args dict, and returns ``(frame, response)``.
-        postprocess : Callable, optional
-            Finalizer applied to the combined result on successful
-            completion (see :data:`_PostProcess`). ``None`` (default)
-            returns the raw combined ``(frame, response)``.
-
-        Returns
-        -------
-        df : pandas.DataFrame
-            Combined data from every successful sub-request.
-        response : httpx.Response or BaseMetadata
-            The aggregated response, or the finalizer's output when
-            ``postprocess`` is supplied.
-
-        Raises
-        ------
-        ChunkInterrupted
-            On a mid-stream transient failure
-            (:class:`QuotaExhausted` for 429,
-            :class:`ServiceInterrupted` for 5xx). The resumable handle
-            is on ``exc.call``.
-        """
-        return ChunkedCall(self, fetch_once, postprocess).resume()
-
 
 def _classify_chunk_error(
     exc: BaseException,
@@ -908,12 +874,11 @@ def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame:
     """
     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 a downstream ``pd.concat([result, geo_page])`` to a
-        # plain DataFrame and strip geometry/CRS.
+        # Preserve the input frame's type (GeoDataFrame vs DataFrame)
+        # even when every chunk is empty: ``_get_resp_data`` returns a
+        # ``gpd.GeoDataFrame()`` for empty geopandas responses, and
+        # returning a bare ``pd.DataFrame()`` here would strip
+        # geometry/CRS for callers that concat the result with geo frames.
         return frames[0] if frames else pd.DataFrame()
     if len(non_empty) == 1:
         # Single-completed-chunk fast path. Return a copy so callers
@@ -1005,8 +970,8 @@ class ChunkedCall:
     Holds the in-flight state (per-sub-request frames and responses)
     and exposes a single :meth:`resume` entry point that drives the
     call from wherever it is to completion — used both for the first
-    invocation (from :meth:`ChunkPlan.execute`) and for subsequent
-    retries after a :class:`ChunkInterrupted`.
+    invocation and for subsequent retries after a
+    :class:`ChunkInterrupted`.
 
     A ``ChunkedCall`` is created internally when a :class:`ChunkPlan`
     executes; callers reach it via :attr:`ChunkInterrupted.call` on
@@ -1054,13 +1019,10 @@ def __init__(
     ) -> None:
         self.plan = plan
         self.fetch_once = fetch_once
-        # Optional finalizer (see ``_PostProcess``), bound here at
-        # construction by the caller that owns post-processing. It rides
-        # on this persistent call object, so it survives repeated
+        # Optional finalizer (see ``_PostProcess``), bound at construction
+        # so it rides on this persistent call object and survives repeated
         # interruptions — every ``ChunkInterrupted`` re-exposes the same
-        # call with its finalizer intact — and ``resume()`` applies it
-        # exactly once, on successful completion. ``None`` returns the raw
-        # combined ``(frame, response)``.
+        # call with its finalizer intact, applied once on success.
         self.postprocess = postprocess
         # Completed (frame, response) pairs keyed by sub-args index;
         # ``resume()`` skips indices already present.
@@ -1172,13 +1134,11 @@ def resume(self) -> tuple[pd.DataFrame, Any]:
                 if reporter is not None:
                     reporter.start_chunk(i + 1)
                 self._issue(i, sub_args)
-            ordered = self._ordered_chunks()
-            frame = _combine_chunk_frames([f for f, _ in ordered])
-            response = _combine_chunk_responses(
-                [r for _, r in ordered], self.plan.canonical_url
-            )
+            # The combined result is exactly what the live ``partial_*``
+            # accessors expose — every sub-request has completed here.
+            frame, response = self.partial_frame, self.partial_response
             # Post-process ONLY on successful completion (and only when a
-            # closure was supplied). The partial/error path keeps the raw
+            # finalizer was bound). The partial/error path keeps the raw
             # combined pair, so an interruption can never be masked by a
             # failure inside post-processing.
             if self.postprocess is None:
@@ -1221,7 +1181,7 @@ def multi_value_chunked(
     *,
     build_request: Callable[..., httpx.Request],
     url_limit: int | None = None,
-) -> Callable[[_FetchOnce], _FetchOnce]:
+) -> Callable[[_FetchOnce], _ChunkedFetch]:
     """
     Decorate a fetch function to transparently chunk over-budget requests.
 
@@ -1265,7 +1225,7 @@ def multi_value_chunked(
     ChunkedCall : Per-sub-request execution and resume semantics.
     """
 
-    def decorator(fetch_once: _FetchOnce) -> _FetchOnce:
+    def decorator(fetch_once: _FetchOnce) -> _ChunkedFetch:
         @functools.wraps(fetch_once)
         def wrapper(
             args: dict[str, Any],
@@ -1274,7 +1234,7 @@ def wrapper(
         ) -> tuple[pd.DataFrame, Any]:
             limit = _WATERDATA_URL_BYTE_LIMIT if url_limit is None else url_limit
             plan = ChunkPlan(args, build_request, limit)
-            return plan.execute(fetch_once, postprocess)
+            return ChunkedCall(plan, fetch_once, postprocess).resume()
 
         return wrapper
 

dataretrieval/waterdata/utils.py

@@ -1294,14 +1294,9 @@ def _postprocess(
     ) -> tuple[pd.DataFrame, BaseMetadata]:
         """Coerce dtypes, arrange/rename columns, sort rows, and wrap the
         response as ``BaseMetadata`` — the shared tail every OGC result
-        passes through.
-
-        Handed to the chunker as the call's finalizer, so it runs once in
-        ``ChunkedCall.resume()`` on successful completion — whether the
-        call ran straight through or was resumed by the caller after a
-        ``ChunkInterrupted``. That single binding keeps a resumed result
-        identical in shape to a non-interrupted one instead of leaking
-        raw, string-typed chunk data with a bare ``httpx.Response``.
+        passes through. Passed to the chunker as the call's finalizer (see
+        ``chunking._PostProcess``) so a resumed result is shaped identically
+        to a non-interrupted one.
         """
         frame = _deal_with_empty(frame, properties, service)
         if convert_type: