refactor(waterdata): Tighten the joint chunker

Three small extractions and one minor optimization. No behavior change;
130 chunker/filter tests stay green.

_iter_sub_args generator yields per-sub-request args dicts; the wrapper's
nested-loop-with-manual-counter collapses to ``for i, sub_args in
enumerate(...)``. The "is this the last sub-request" branch in the
quota-floor check flips to ``if i == total - 1: continue`` so the gate
is a guard clause rather than the body of an inverse condition.

_finalize_response folds the ``_combine_chunk_responses(responses);
response.url = canonical_url`` pattern (used in both the success path
and the QuotaExhausted partial-state payload) into one helper.

_filter_candidates generator emits ``(filter_chunks, worst_filter)``
pairs for each candidate filter chunk count; ``_plan_joint`` then
iterates candidates uniformly without the ``if filter_chunkable: ...
else: ...`` fork. The redundant ``filter_chunkable`` flag is gone —
``len(clauses) >= 2`` is the single truth.

Per-iteration optimization: ``{**args, **list_overrides}`` was being
recomputed for every filter chunk; now built once per outer combo and
reused (or shallow-overridden when a filter substitution applies).

Module constants ``_LIST_SEP = ","`` and ``_OR_SEP = " OR "`` replace
the scattered string literals — the two chunking dimensions are now
self-documenting at every call site that sizes them.

Author: thodson-usgs

dataretrieval/waterdata/chunking.py

@@ -34,7 +34,7 @@
 import functools
 import itertools
 import math
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from typing import Any
 from urllib.parse import quote_plus
 
@@ -98,6 +98,14 @@
 # so a missing/malformed header doesn't trigger spurious aborts.
 _QUOTA_UNKNOWN = 10**9
 
+# 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
+# (``filter=a='1' OR a='2'``). Pinned as constants so the URL-byte
+# sizing helper and the partition logic agree on the join shape.
+_LIST_SEP = ","
+_OR_SEP = " OR "
+
 
 class RequestTooLarge(ValueError):
     """Raised when a chunked request cannot be issued. Either the URL
@@ -214,7 +222,7 @@ def _worst_case_args(
     in by the caller)."""
     out = dict(base_args)
     for k, chunks in list_plan.items():
-        out[k] = max(chunks, key=lambda c: _joined_url_bytes(c, ","))
+        out[k] = max(chunks, key=lambda c: _joined_url_bytes(c, _LIST_SEP))
     return out
 
 
@@ -282,7 +290,9 @@ def _plan_list_chunks(
             if len(chunk) > 1
         )
         biggest = max(
-            splittable, key=lambda t: _joined_url_bytes(t[2], ","), default=None
+            splittable,
+            key=lambda t: _joined_url_bytes(t[2], _LIST_SEP),
+            default=None,
         )
         if biggest is None:
             raise RequestTooLarge(
@@ -320,6 +330,32 @@ def _filter_chunk_counts(n_clauses: int) -> list[int]:
     return counts
 
 
+def _filter_candidates(
+    clauses: list[str], original_filter: str | None
+) -> Iterator[tuple[list[str | None], str | None]]:
+    """Yield ``(filter_chunks, worst_filter)`` for each candidate filter
+    chunk count. ``filter_chunks`` is the list of OR-joined sub-filters
+    the wrapper will substitute one at a time; ``worst_filter`` is the
+    longest URL-encoded chunk, used by the planner to size the list dims
+    against the most demanding sub-request.
+
+    Falls through to a single ``(filter_chunks=[original_filter], None)``
+    candidate when the filter has no top-level OR splits (single clause,
+    cql-json, missing filter): the wrapper still iterates that one
+    "chunk" but the planner can skip substituting a filter into the
+    URL probe."""
+    if len(clauses) < 2:
+        chunks: list[str | None] = (
+            [original_filter] if original_filter is not None else [None]
+        )
+        yield chunks, None
+        return
+    for k in _filter_chunk_counts(len(clauses)):
+        groups = _partition_clauses(clauses, k)
+        worst = max(groups, key=lambda g: _joined_url_bytes(g, _OR_SEP))
+        yield [_OR_SEP.join(g) for g in groups], _OR_SEP.join(worst)
+
+
 def _plan_joint(
     args: dict[str, Any],
     build_request: Callable[..., Any],
@@ -358,42 +394,30 @@ def _plan_joint(
     if _is_chunkable(filter_expr, args.get("filter_lang")):
         _check_numeric_filter_pitfall(filter_expr)
         clauses = _split_top_level_or(filter_expr)
-    # Filter is chunkable only when there are ≥2 top-level OR clauses;
-    # a single clause can't be split losslessly.
-    filter_chunkable = len(clauses) >= 2
 
-    if not _chunkable_params(args) and not filter_chunkable:
+    if not _chunkable_params(args) and len(clauses) < 2:
         return None
     if _request_bytes(build_request(**args)) <= url_limit:
         return None
 
-    candidate_counts = _filter_chunk_counts(len(clauses)) if filter_chunkable else [1]
     best: tuple[int, dict[str, list[list[Any]]], list[str | None]] | None = None
     last_error: RequestTooLarge | None = None
 
-    for k in candidate_counts:
-        if filter_chunkable:
-            groups = _partition_clauses(clauses, k)
-            worst_group = max(groups, key=lambda g: _joined_url_bytes(g, " OR "))
-            filter_chunks: list[str | None] = [" OR ".join(g) for g in groups]
-            plan_args = {**args, "filter": " OR ".join(worst_group)}
-        else:
-            filter_chunks = [filter_expr] if filter_expr is not None else [None]
-            plan_args = args
-
-        per_filter_cap = max(1, max_chunks // k)
+    for filter_chunks, worst_filter in _filter_candidates(clauses, filter_expr):
+        k = len(filter_chunks)
+        plan_args = args if worst_filter is None else {**args, "filter": worst_filter}
         try:
             list_plan = _plan_list_chunks(
-                plan_args, build_request, url_limit, per_filter_cap
+                plan_args, build_request, url_limit, max(1, max_chunks // k)
             )
         except RequestTooLarge as exc:
             last_error = exc
             continue
         if list_plan is None:
             list_plan = {}
-        # When there are no list dims to shrink, ``_plan_list_chunks``
-        # returns ``None`` regardless of whether the request actually
-        # fits. Filter chunking alone has to close the gap — verify it.
+        # ``_plan_list_chunks`` returns ``None`` when no list dims are
+        # chunkable, regardless of whether the request actually fits.
+        # Filter chunking alone has to close the gap — verify it.
         if not list_plan and _request_bytes(build_request(**plan_args)) > url_limit:
             continue
         total = _plan_total(list_plan, k)
@@ -425,6 +449,34 @@ def _read_remaining(response: requests.Response) -> int:
         return _QUOTA_UNKNOWN
 
 
+def _iter_sub_args(
+    args: dict[str, Any],
+    list_plan: dict[str, list[list[Any]]],
+    filter_chunks: list[str | None],
+) -> Iterator[dict[str, Any]]:
+    """Yield the substituted ``args`` for each sub-request in the joint
+    plan, in deterministic order: list-dim cartesian product (insertion
+    order, Python 3.7+ guarantee) crossed with filter chunks."""
+    list_keys = list(list_plan)
+    list_combos = (
+        itertools.product(*(list_plan[k] for k in list_keys)) if list_plan else [()]
+    )
+    for combo in list_combos:
+        base = {**args, **dict(zip(list_keys, combo))}
+        for filter_chunk in filter_chunks:
+            yield base if filter_chunk is None else {**base, "filter": filter_chunk}
+
+
+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],
@@ -481,44 +533,32 @@ def wrapper(
 
             list_plan, filter_chunks = plan
             canonical_url = build_request(**args).url
-
-            list_keys = list(list_plan)
-            list_combos = (
-                list(itertools.product(*(list_plan[k] for k in list_keys)))
-                if list_plan
-                else [()]
-            )
-            total = len(list_combos) * len(filter_chunks)
+            total = _plan_total(list_plan, len(filter_chunks))
 
             frames: list[pd.DataFrame] = []
             responses: list[requests.Response] = []
-            i = 0
-            for combo in list_combos:
-                for filter_chunk in filter_chunks:
-                    sub_args = dict(args)
-                    sub_args.update(zip(list_keys, combo))
-                    if filter_chunk is not None:
-                        sub_args["filter"] = filter_chunk
-                    frame, response = fetch_once(sub_args)
-                    frames.append(frame)
-                    responses.append(response)
-                    if i < total - 1:
-                        remaining = _read_remaining(response)
-                        if remaining < floor:
-                            partial = _combine_chunk_responses(responses)
-                            partial.url = canonical_url
-                            raise QuotaExhausted(
-                                partial_frame=_combine_chunk_frames(frames),
-                                partial_response=partial,
-                                completed_chunks=i + 1,
-                                total_chunks=total,
-                                remaining=remaining,
-                            )
-                    i += 1
-
-            combined = _combine_chunk_responses(responses)
-            combined.url = canonical_url
-            return _combine_chunk_frames(frames), combined
+            for i, sub_args in enumerate(
+                _iter_sub_args(args, list_plan, filter_chunks)
+            ):
+                frame, response = fetch_once(sub_args)
+                frames.append(frame)
+                responses.append(response)
+                if i == total - 1:
+                    continue  # last chunk; no next sub-request to gate
+                remaining = _read_remaining(response)
+                if remaining < floor:
+                    raise QuotaExhausted(
+                        partial_frame=_combine_chunk_frames(frames),
+                        partial_response=_finalize_response(responses, canonical_url),
+                        completed_chunks=i + 1,
+                        total_chunks=total,
+                        remaining=remaining,
+                    )
+
+            return (
+                _combine_chunk_frames(frames),
+                _finalize_response(responses, canonical_url),
+            )
 
         return wrapper  # type: ignore[return-value]