refactor: replace ChunkInterrupted snapshot copies with live-view properties

Drop the .copy() snapshot of call.partial_frame and the stored
partial_response from ChunkInterrupted.__init__. Instead, partial_frame
and partial_response are now @Property delegates to self.call when
available, falling back to empty/None when call is None.

This avoids copying potentially large DataFrames on every
ChunkInterrupted raise — the primary usage pattern is exc.call.resume(),
not inspecting the partial data on the exception itself.

Pickling still works: __getstate__ snapshots the current partial_frame
and partial_response into private backing attributes that __setstate__
restores, so the degraded (call=None) unpickled form retains the data
that was available at serialization time.

Author: thodson-usgs

dataretrieval/ogc/interruptions.py

@@ -65,14 +65,14 @@ class ChunkInterrupted(DataRetrievalError):
     total_chunks : int
         Total sub-requests in the plan.
     partial_frame : pandas.DataFrame
-        Combined frame of work completed by the moment this exception
-        was raised. Snapshot at raise time — does NOT advance on a
-        later ``call.resume()`` (use ``exc.call.partial_frame`` for
-        the live view).
+        Live view of work completed so far — delegates to
+        ``call.partial_frame``. Advances as ``call.resume()`` makes
+        further progress. Returns an empty ``DataFrame`` when
+        ``call`` is ``None`` (degraded/unpickled state).
     partial_response : httpx.Response or None
-        Raw aggregate response covering the completed sub-requests at
-        raise time; ``None`` if nothing had completed yet. Same snapshot
-        semantics as ``partial_frame``. (Raw, not finalized — use
+        Live view of the raw aggregate response — delegates to
+        ``call.partial_response``. ``None`` when ``call`` is ``None``
+        or nothing has completed yet. (Raw, not finalized — use
         ``exc.call.resume()`` for the finalized ``(df, metadata)`` result.)
 
     Examples
@@ -128,16 +128,32 @@ 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.
-        if call is None:
-            self.partial_frame: pd.DataFrame = pd.DataFrame()
-            self.partial_response: httpx.Response | None = None
-        else:
-            self.partial_frame = call.partial_frame.copy()
-            self.partial_response = call.partial_response
+
+    @property
+    def partial_frame(self) -> pd.DataFrame:
+        """Live view of work completed so far — delegates to ``call.partial_frame``.
+
+        Returns an empty ``DataFrame`` when ``call`` is ``None`` (degraded /
+        unpickled state), unless a pickled snapshot was restored via
+        ``__setstate__``.
+        """
+        if self.call is not None:
+            return self.call.partial_frame
+        # Degraded / unpickled: return the snapshot captured at pickle time,
+        # or an empty frame if no snapshot exists (hand-constructed exc).
+        return getattr(self, "_pickled_partial_frame", pd.DataFrame())
+
+    @property
+    def partial_response(self) -> httpx.Response | None:
+        """Live view of the raw aggregate response.
+
+        Delegates to ``call.partial_response``.
+        Returns ``None`` when ``call`` is ``None`` (degraded / unpickled state)
+        and no pickled snapshot exists.
+        """
+        if self.call is not None:
+            return self.call.partial_response
+        return getattr(self, "_pickled_partial_response", None)
 
     def __getstate__(self) -> dict[str, Any]:
         # Drop the live ChunkedCall before pickling: its ``.fetch`` is an
@@ -146,7 +162,25 @@ def __getstate__(self) -> dict[str, Any]:
         # 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).
-        return {**super().__getstate__(), "call": None}
+        #
+        # Since ``partial_frame`` and ``partial_response`` are now live
+        # properties (delegating to ``self.call``), we must snapshot them
+        # here so the pickled state includes the data that was available
+        # at serialization time.
+        state = {**super().__getstate__(), "call": None}
+        state["_pickled_partial_frame"] = self.partial_frame.copy()
+        state["_pickled_partial_response"] = self.partial_response
+        return state
+
+    def __setstate__(self, state: dict[str, Any] | None) -> None:
+        state = state or {}
+        # Restore the pickled partial snapshots into private backing attrs
+        # that the properties will fall back to when call is None.
+        self._pickled_partial_frame = state.pop(
+            "_pickled_partial_frame", pd.DataFrame()
+        )
+        self._pickled_partial_response = state.pop("_pickled_partial_response", None)
+        super().__setstate__(state)
 
 
 class QuotaExhausted(ChunkInterrupted):

tests/waterdata_chunking_test.py

@@ -954,13 +954,11 @@ async def fetch(args):
     assert err.partial_response is not None
 
 
-def test_partial_frame_snapshot_stable_across_resume():
-    """``exc.partial_frame`` / ``exc.partial_response`` snapshot the
-    state at raise time. Calling ``exc.call.resume()`` advances the
-    underlying ``ChunkedCall`` but must NOT mutate the snapshot on
-    the exception — otherwise a diagnostic that reads
-    ``exc.partial_frame`` after a resume sees post-resume state under
-    a name that promises pre-resume state."""
+def test_partial_frame_is_live_view_across_resume():
+    """``exc.partial_frame`` is now a live view delegating to
+    ``exc.call.partial_frame``. Calling ``exc.call.resume()`` advances the
+    underlying ``ChunkedCall`` and ``exc.partial_frame`` reflects that — it
+    is no longer a frozen snapshot."""
     state = {"i": 0, "blow_up": True}
 
     async def fetch(args):
@@ -977,27 +975,24 @@ async def fetch(args):
     with pytest.raises(QuotaExhausted) as excinfo:
         decorated({"sites": ["S1" * 10, "S2" * 10, "S3" * 10, "S4" * 10, "S5" * 10]})
     err = excinfo.value
-    snapshot_rows = len(err.partial_frame)
-    assert snapshot_rows > 0  # two chunks worth of data captured
+    pre_resume_rows = len(err.partial_frame)
+    assert pre_resume_rows > 0  # two chunks worth of data captured
 
-    # Resume; the live view on .call grows.
+    # Resume; the live view on .call grows — and so does exc.partial_frame.
     state["blow_up"] = False
     err.call.resume()
-    assert len(err.call.partial_frame) > snapshot_rows
+    assert len(err.call.partial_frame) > pre_resume_rows
 
-    # The exception's snapshot must NOT advance.
-    assert len(err.partial_frame) == snapshot_rows
+    # The exception's partial_frame IS the live view — it advances too.
+    assert len(err.partial_frame) == len(err.call.partial_frame)
 
 
-def test_partial_frame_snapshot_is_a_copy_when_single_chunk():
-    """``_combine_chunk_frames`` returns ``non_empty[0]`` verbatim on
-    its single-frame fast path. ``ChunkInterrupted.__init__`` must
-    therefore defensively ``.copy()`` so an in-place mutation of the
-    underlying chunk frame (e.g. user diagnostic code adding a
-    column on the live view) doesn't leak through the snapshot.
-    Companion to ``test_partial_frame_snapshot_stable_across_resume``,
-    which uses ≥2 completed chunks and so goes through
-    ``pd.concat`` (which already produces a fresh frame)."""
+def test_partial_frame_is_live_view_when_single_chunk():
+    """``partial_frame`` is a live view delegating to ``call.partial_frame``.
+    When a single chunk completes, ``partial_frame`` still reflects the
+    current state of ``call.partial_frame`` (which recomputes on each
+    access). This verifies the live-view semantics even on the single-
+    chunk fast path."""
     state = {"i": 0, "blow_up": True}
 
     async def fetch(args):
@@ -1020,12 +1015,11 @@ async def fetch(args):
     err = excinfo.value
     assert err.completed_chunks == 1
 
-    snapshot_cols = list(err.partial_frame.columns)
-    # Mutate the underlying chunk in place — the snapshot must NOT
-    # reflect the mutation.
-    err.call._chunks[0][0]["extra"] = 0
-    assert list(err.partial_frame.columns) == snapshot_cols
-    assert "extra" not in err.partial_frame.columns
+    # The live view delegates to call.partial_frame; verify it's
+    # accessible and reflects the call's current state.
+    assert not err.partial_frame.empty
+    assert list(err.partial_frame.columns) == ["sites"]
+    assert err.partial_frame.equals(err.call.partial_frame)
 
 
 def test_combine_chunk_responses_returns_independent_headers():