diff --git a/python/lib/sift_client/_internal/low_level_wrappers/data.py b/python/lib/sift_client/_internal/low_level_wrappers/data.py index e16cc750c..678fc7eb6 100644 --- a/python/lib/sift_client/_internal/low_level_wrappers/data.py +++ b/python/lib/sift_client/_internal/low_level_wrappers/data.py @@ -2,8 +2,9 @@ import asyncio import logging +import time from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Tuple, cast +from typing import TYPE_CHECKING, Any, Callable, Tuple, cast import pandas as pd from pydantic import BaseModel, ConfigDict @@ -19,6 +20,7 @@ from sift_client._internal.disk_cache import DiskCache from sift_client._internal.low_level_wrappers.base import LowLevelClientBase from sift_client._internal.time import to_timestamp_nanos +from sift_client._internal.util.progress import alive_bar from sift_client.sift_types.channel import Channel, ChannelDataType from sift_client.transport import WithGrpcClient @@ -589,6 +591,52 @@ async def _get_data_impl( response = cast("GetDataResponse", response) return response.data, response.next_page_token # type: ignore # mypy doesn't know RepeatedCompositeFieldContainer can be treated like a list + async def _paginate_channel_data( + self, + *, + kwargs: dict[str, Any], + page_size: int | None, + max_points: int | None, + on_page: Callable[[int], None] | None = None, + ) -> list[dict[str, pd.DataFrame]]: + """Page one ``GetData`` query up to a point budget. + + ``max_points`` (the caller's ``limit``) caps points, but each page + is one proto holding up to ``page_size`` points, so the budget + drives ``page_size`` and the page count rather than a proto count. + Deserialized once here so :meth:`_merge_pages` needn't parse again. + """ + if max_points == 0: + return [] + + if max_points is not None: + # Never pull a larger page than the whole budget; then cover + # the budget in ceil(budget / page_size) equal pages. + if page_size is None or page_size > max_points: + page_size = max_points + pages_remaining: int | None = -(-max_points // page_size) # ceil div + else: + page_size = page_size or CHANNELS_DEFAULT_PAGE_SIZE + pages_remaining = None # exhaust via the page token + + deserialized: list[dict[str, pd.DataFrame]] = [] + page_token: str | None = "" + while True: + page, page_token = await self._get_data_impl( + page_size=page_size, page_token=page_token, **kwargs + ) + frames = [self.try_deserialize_channel_data(d) for d in page] + deserialized.extend(frames) + if on_page is not None: + on_page(sum(len(df) for frame in frames for df in frame.values())) + if pages_remaining is not None: + pages_remaining -= 1 + if pages_remaining <= 0: + break + if not page_token: + break + return deserialized + def _update_cache( self, *, @@ -724,6 +772,7 @@ async def get_channel_data( max_results: int | None = None, page_size: int | None = None, ignore_cache: bool = False, + show_progress: bool = False, ) -> dict[str, pd.DataFrame]: """Get the data for a channel during a run.""" ret_data: dict[str, pd.DataFrame] = {} @@ -743,6 +792,11 @@ async def get_channel_data( fully_uncached: list[str] = [] partial_gaps: list[tuple[str, list[TimeRange]]] = [] fetched_ranges_per_channel: dict[str, list[TimeRange]] = {} + # Points served from cache, counted the same way wire pages are + # (per element frame, so ``.size`` covers bitfield columns). Seeds + # the progress total so the receipt reflects the full result, but + # stays out of the fetched-throughput rate. + cached_points = 0 for channel in channels: cid = channel.id_ @@ -754,6 +808,7 @@ async def get_channel_data( cached_data, gaps = self.channel_cache.get_range(cid, run_id, start_time, end_time) if cached_data is not None: + cached_points += int(cached_data.size) # Slice per column so each result key carries only its # own element frame (matches the per-element shape # ``try_deserialize_channel_data`` produces; without @@ -770,46 +825,133 @@ async def get_channel_data( else: partial_gaps.append((cid, gaps)) - tasks = [] + # One spec per wire query: the fetch kwargs plus a human-readable + # label for the progress bar. Built before the bar so its total + # is the dispatched-query count, then turned into coroutines + # inside the bar block where the ``on_page`` callback exists. + specs: list[tuple[dict[str, Any], str]] = [] + id_to_name = {channel.id_: channel.name for channel in channels} # Batch fully-uncached channels (sharing the full requested # range) into one wire call each. batch_size = REQUEST_BATCH_SIZE for i in range(0, len(fully_uncached), batch_size): batch = fully_uncached[i : i + batch_size] - task = asyncio.create_task( - self._handle_pagination( - self._get_data_impl, - kwargs={ + specs.append( + ( + { "channel_ids": batch, "run_id": run_id, "start_time": start_time, "end_time": end_time, }, - page_size=page_size, - max_results=max_results, + ", ".join(id_to_name.get(cid, cid) for cid in batch), ) ) - tasks.append(task) # Partial gaps: one fetch per (channel, gap). for cid, gaps in partial_gaps: for gap_start, gap_end in gaps: - task = asyncio.create_task( - self._handle_pagination( - self._get_data_impl, - kwargs={ + specs.append( + ( + { "channel_ids": [cid], "run_id": run_id, "start_time": gap_start, "end_time": gap_end, }, - page_size=page_size, - max_results=max_results, + f"{id_to_name.get(cid, cid)} [{gap_start:%H:%M:%S}-{gap_end:%H:%M:%S}]", ) ) - tasks.append(task) - pages = await asyncio.gather(*tasks) + with alive_bar( + len(specs), + title="Fetching channel data", + stats=False, # hide the built-in channels/s rate; we show points/s below + dual_line=True, # channel fill on line one, points + rate on line two + receipt_text=True, # keep the final points + rate line as a receipt + disable=not show_progress, + ) as bar: + # The fill tracks channels done; points and throughput are + # tracked here and rendered on the second line, since points + # have no clean total to fill against. Fetches run + # concurrently, so name the set in flight as a group rather + # than a single "current" channel. Total points include the + # cache; the rate is fetched-only, since cache reads have no + # meaningful throughput and would distort pts/s. + in_flight: list[str] = [] + fetched_points = 0 + start = time.monotonic() + + def _render() -> None: + total = cached_points + fetched_points + parts = [f"{total:,} pts"] + # Key the split on points actually returned from the wire, + # not on whether a fetch was dispatched: a warm call can + # still issue an empty leading-gap query, which should + # read as "all cached", not "0 fetched". + if fetched_points and cached_points: + parts[0] += f" ({fetched_points:,} fetched · {cached_points:,} cached)" + elif cached_points: + parts[0] += " (all cached)" + if fetched_points: + elapsed = time.monotonic() - start + rate = fetched_points / elapsed if elapsed > 0 else 0.0 + parts.append(f"{rate:,.0f} pts/s") + if in_flight: + shown = in_flight[:3] + names = ", ".join(shown) + extra = len(in_flight) - len(shown) + if extra: + names += f" (+{extra} more)" + parts.append(names) + bar.text(" · ".join(parts)) + + def on_page(points: int) -> None: + nonlocal fetched_points + fetched_points += points + _render() + + async def _tick(kwargs: dict[str, Any], label: str) -> Any: + in_flight.append(label) + _render() + try: + return await self._paginate_channel_data( + kwargs=kwargs, + page_size=page_size, + max_points=max_results, + on_page=on_page, + ) + finally: + in_flight.remove(label) + _render() + bar() + + # Seed the line up front so a fully-cached call (no fetch + # tasks) still shows its cached-point summary in the receipt. + _render() + # Keep ``gather`` (not ``as_completed``) so pages stay in + # submission order: ``_merge_pages`` lets later-positioned + # pages win on duplicate timestamps. + pages = await asyncio.gather(*[_tick(kw, lbl) for kw, lbl in specs]) + + # Replace the live line with a past-tense summary so the + # receipt reads as a result, not a mid-flight tick. The rate + # is labelled as an average and covers only fetched points. + total = cached_points + fetched_points + elapsed = time.monotonic() - start + avg = fetched_points / elapsed if elapsed > 0 else 0.0 + if fetched_points and cached_points: + summary = ( + f"Loaded {total:,} points: {fetched_points:,} fetched " + f"(avg {avg:,.0f}/s), {cached_points:,} cached" + ) + elif fetched_points: + summary = f"Loaded {total:,} points (avg {avg:,.0f}/s)" + elif cached_points: + summary = f"Loaded {total:,} points from cache" + else: + summary = f"Loaded {total:,} points" + bar.text(summary) ret_data = self._merge_pages(pages, initial=ret_data) # Pure cache hits never reach ``_update_cache`` because nothing @@ -830,12 +972,17 @@ async def get_channel_data( def _merge_pages( self, - pages: list[list[Any]], + pages: list[list[dict[str, pd.DataFrame]]], *, initial: dict[str, pd.DataFrame], ) -> dict[str, pd.DataFrame]: """Flatten paged channel data + any cached slices into one DataFrame per channel. + ``pages`` is already deserialized: one entry per fetch task, each + a list of per-page ``{name: frame}`` dicts from + :meth:`_paginate_channel_data`. Parsing happens there so the wire + payload is read once. + ``initial`` carries the cached slices ``get_channel_data`` stitched inline via :meth:`ChannelDataCache.get_range` before dispatching wire fetches for the gaps. Cached entries are @@ -846,8 +993,8 @@ def _merge_pages( """ per_channel_frames: dict[str, list[pd.DataFrame]] = {} for page in pages: - for data in page: - for name, df in self.try_deserialize_channel_data(data).items(): + for frame in page: + for name, df in frame.items(): per_channel_frames.setdefault(name, []).append(df) ret_data: dict[str, pd.DataFrame] = dict(initial) diff --git a/python/lib/sift_client/_internal/util/file.py b/python/lib/sift_client/_internal/util/file.py index ac1b83f57..14351c15f 100644 --- a/python/lib/sift_client/_internal/util/file.py +++ b/python/lib/sift_client/_internal/util/file.py @@ -132,13 +132,21 @@ def download_file( return output_path -def extract_zip(zip_path: Path, output_dir: Path, *, delete_zip: bool = True) -> list[Path]: +def extract_zip( + zip_path: Path, + output_dir: Path, + *, + delete_zip: bool = True, + show_progress: bool = False, +) -> list[Path]: """Extract a zip file to a directory. Args: zip_path: Path to the zip file. output_dir: Directory to extract contents into. Created if it doesn't exist. delete_zip: If True (default), delete the zip file after extraction. + show_progress: If True, display a progress bar naming each file as it is + extracted. Defaults to False. Returns: List of paths to the extracted files (excludes directories). @@ -149,7 +157,11 @@ def extract_zip(zip_path: Path, output_dir: Path, *, delete_zip: bool = True) -> output_dir.mkdir(parents=True, exist_ok=True) with zipfile.ZipFile(zip_path, "r") as zip_file: names = zip_file.namelist() - zip_file.extractall(output_dir) + with alive_bar(len(names), title="Extracting", disable=not show_progress) as bar: + for name in names: + bar.text(name) + zip_file.extract(name, output_dir) + bar() if delete_zip: try: zip_path.unlink() diff --git a/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_data.py b/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_data.py index ba5a8f98c..31fb8cd7b 100644 --- a/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_data.py +++ b/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_data.py @@ -37,7 +37,7 @@ from contextlib import contextmanager from datetime import datetime, timedelta, timezone -from typing import Any, Iterator, cast +from typing import Any, Iterator, Mapping, Sequence, cast from unittest.mock import MagicMock, patch import pandas as pd @@ -134,8 +134,9 @@ def _client_with_cache(tmp_path, subdir: str = "cache") -> DataLowLevelClient: def _patch_deserializer(sentinel_to_frames: dict[str, dict[str, pd.DataFrame]]) -> Any: """Patch ``try_deserialize_channel_data`` to translate string sentinels. - Lets tests pass strings in lieu of protos. Returned object is a context - manager; callers use ``with _patch_deserializer(...):``. + Lets ``_fake_grpc`` pass strings in lieu of protos; the paginator + deserializes each page through this patched reader. Returned object is + a context manager; callers use ``with _patch_deserializer(...):``. """ return patch.object( DataLowLevelClient, @@ -147,7 +148,7 @@ def _patch_deserializer(sentinel_to_frames: dict[str, dict[str, pd.DataFrame]]) @contextmanager def _fake_grpc( client: DataLowLevelClient, - channel_to_pages: dict[str, list[pd.DataFrame | dict[str, pd.DataFrame]]], + channel_to_pages: Mapping[str, Sequence[pd.DataFrame | dict[str, pd.DataFrame]]], ) -> Iterator[list[dict[str, Any]]]: """Mock the gRPC boundary so each "page" is a sentinel string. @@ -957,16 +958,14 @@ def test_no_fresh_data_returns_initial(self, pages: list) -> None: """No fresh pages → initial dict passes through by identity.""" client = DataLowLevelClient(MagicMock()) initial_df = _frame("chan", rows=5) - with _patch_deserializer({}): - result = client._merge_pages(pages=pages, initial={"chan": initial_df}) + result = client._merge_pages(pages=pages, initial={"chan": initial_df}) assert result["chan"] is initial_df def test_single_frame_skips_concat(self) -> None: """One frame for a channel → returned by identity, no concat call.""" only_df = _frame("chan", rows=5) client = DataLowLevelClient(MagicMock()) - with _patch_deserializer({"p1": {"chan": only_df}}): - result = client._merge_pages(pages=[["p1"]], initial={}) + result = client._merge_pages(pages=[[{"chan": only_df}]], initial={}) assert result["chan"] is only_df def test_disjoint_pages_concat_in_order(self) -> None: @@ -975,9 +974,9 @@ def test_disjoint_pages_concat_in_order(self) -> None: df2 = _frame("chan", rows=10, start=_NOW + timedelta(minutes=1), offset=10, freq="s") df3 = _frame("chan", rows=10, start=_NOW + timedelta(minutes=2), offset=20, freq="s") client = DataLowLevelClient(MagicMock()) - sentinels = {"p1": {"chan": df1}, "p2": {"chan": df2}, "p3": {"chan": df3}} - with _patch_deserializer(sentinels): - result = client._merge_pages(pages=[["p1", "p2"], ["p3"]], initial={}) + result = client._merge_pages( + pages=[[{"chan": df1}, {"chan": df2}], [{"chan": df3}]], initial={} + ) expected = pd.concat([df1, df2, df3]).groupby(level=0).last() pd.testing.assert_frame_equal(result["chan"].sort_index(), expected.sort_index()) assert len(result["chan"]) == 30 @@ -992,8 +991,7 @@ def test_overlapping_timestamps_later_page_wins(self) -> None: df_first = pd.DataFrame({"chan": [0] * 5}, index=index) df_second = pd.DataFrame({"chan": [99] * 5}, index=index) client = DataLowLevelClient(MagicMock()) - with _patch_deserializer({"p1": {"chan": df_first}, "p2": {"chan": df_second}}): - result = client._merge_pages(pages=[["p1", "p2"]], initial={}) + result = client._merge_pages(pages=[[{"chan": df_first}, {"chan": df_second}]], initial={}) assert (result["chan"]["chan"] == 99).all() def test_cached_slice_folded_in_first_and_loses_on_overlap(self) -> None: @@ -1006,8 +1004,7 @@ def test_cached_slice_folded_in_first_and_loses_on_overlap(self) -> None: cached = pd.DataFrame({"chan": [-1] * 5}, index=index) fresh = pd.DataFrame({"chan": [42] * 5}, index=index) client = DataLowLevelClient(MagicMock()) - with _patch_deserializer({"p1": {"chan": fresh}}): - result = client._merge_pages(pages=[["p1"]], initial={"chan": cached}) + result = client._merge_pages(pages=[[{"chan": fresh}]], initial={"chan": cached}) assert (result["chan"]["chan"] == 42).all() def test_multiple_channels_independent(self) -> None: @@ -1016,9 +1013,7 @@ def test_multiple_channels_independent(self) -> None: a2 = _frame("a", rows=5, start=_NOW + timedelta(minutes=1), offset=5, freq="s") b1 = _frame("b", rows=5, start=_NOW, offset=100, freq="s") client = DataLowLevelClient(MagicMock()) - sentinels = {"p_a1": {"a": a1}, "p_a2": {"a": a2}, "p_b1": {"b": b1}} - with _patch_deserializer(sentinels): - result = client._merge_pages(pages=[["p_a1", "p_b1"], ["p_a2"]], initial={}) + result = client._merge_pages(pages=[[{"a": a1}, {"b": b1}], [{"a": a2}]], initial={}) assert len(result["a"]) == 10 assert len(result["b"]) == 5 assert (result["b"]["b"] >= 100).all() @@ -1029,8 +1024,7 @@ def test_does_not_mutate_initial(self) -> None: initial = {"chan": cached} fresh = _frame("chan", rows=5, start=_NOW + timedelta(seconds=1), offset=10) client = DataLowLevelClient(MagicMock()) - with _patch_deserializer({"p1": {"chan": fresh}}): - client._merge_pages(pages=[["p1"]], initial=initial) + client._merge_pages(pages=[[{"chan": fresh}]], initial=initial) assert initial["chan"] is cached @@ -1130,6 +1124,89 @@ async def test_multi_page_response_concatenated_per_channel(self) -> None: expected = pd.concat([p1, p2, p3]).groupby(level=0).last() pd.testing.assert_frame_equal(result["c1"].sort_index(), expected.sort_index()) + @pytest.mark.asyncio + async def test_limit_bounds_round_trips_to_one_page(self) -> None: + """A point budget met by a single page stops after one wire call. + + Guards the pagination fix: ``limit`` is a data-point budget that + sizes ``page_size`` and the page count, not a proto count. With + the pre-fix proto-counting loop this paged until the token + emptied (3 calls here) even though one page covers the budget. + """ + client = DataLowLevelClient(MagicMock()) + pages = [_frame("c1", rows=10, start=_NOW + timedelta(seconds=i)) for i in range(3)] + with _fake_grpc(client, {"c1": pages}) as call_log: + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + max_results=100, + page_size=100, + ignore_cache=True, + ) + assert len(call_log) == 1, "budget covered by one page must not keep paging" + + @pytest.mark.asyncio + async def test_limit_pages_ceil_of_budget_over_page_size(self) -> None: + """Round-trips equal ceil(limit / page_size) when the page is smaller.""" + client = DataLowLevelClient(MagicMock()) + pages = [_frame("c1", rows=25, start=_NOW + timedelta(seconds=i)) for i in range(6)] + with _fake_grpc(client, {"c1": pages}) as call_log: + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + max_results=100, + page_size=25, # ceil(100 / 25) == 4 pages + ignore_cache=True, + ) + assert len(call_log) == 4 + + @pytest.mark.asyncio + async def test_no_limit_exhausts_all_pages(self) -> None: + """Without a limit, paging follows the token to exhaustion.""" + client = DataLowLevelClient(MagicMock()) + pages = [_frame("c1", rows=10, start=_NOW + timedelta(seconds=i)) for i in range(3)] + with _fake_grpc(client, {"c1": pages}) as call_log: + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ignore_cache=True, + ) + assert len(call_log) == 3 + + @pytest.mark.asyncio + async def test_limit_above_available_terminates_on_token(self) -> None: + """A budget larger than the data stops when the token empties, not before.""" + client = DataLowLevelClient(MagicMock()) + pages = [_frame("c1", rows=500, start=_NOW + timedelta(seconds=i)) for i in range(2)] + with _fake_grpc(client, {"c1": pages}) as call_log: + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + max_results=100_000, # ceil would ask for more pages than exist + page_size=500, + ignore_cache=True, + ) + assert len(call_log) == 2 + + @pytest.mark.asyncio + async def test_limit_zero_returns_empty_without_wire_calls(self) -> None: + """A zero budget short-circuits before touching the wire.""" + client = DataLowLevelClient(MagicMock()) + with _fake_grpc(client, {"c1": [_frame("c1")]}) as call_log: + result = await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + max_results=0, + ignore_cache=True, + ) + assert call_log == [] + assert result == {} + @pytest.mark.asyncio async def test_cache_hit_short_circuits_grpc(self, tmp_path) -> None: """Second request for the same channel + window skips ``_get_data_impl``. @@ -1617,3 +1694,167 @@ async def test_get_channel_data_cached_bitfield_returns_per_element_frames( ) finally: client.channel_cache.store.close() + + +@contextmanager +def _capture_bar_total() -> Iterator[list[int | None]]: + """Patch ``alive_bar`` in the data module and capture each ``total`` arg. + + ``get_channel_data`` opens exactly one bar per call whose ``total`` is + the number of dispatched wire-fetch tasks, so the captured value is the + task count. + """ + totals: list[int | None] = [] + + @contextmanager + def _fake_alive_bar(total: int | None = None, *_args: Any, **_kwargs: Any) -> Iterator[Any]: + totals.append(total) + yield MagicMock() + + with patch( + "sift_client._internal.low_level_wrappers.data.alive_bar", + _fake_alive_bar, + ): + yield totals + + +@contextmanager +def _capture_bar_text() -> Iterator[list[str]]: + """Patch ``alive_bar`` and capture every ``bar.text(...)`` string. + + The bar's fill tracks channels; points loaded and throughput are + rendered into the situational text, one update per page. + """ + texts: list[str] = [] + + @contextmanager + def _fake_alive_bar(total: int | None = None, *_args: Any, **_kwargs: Any) -> Iterator[Any]: + bar = MagicMock() + bar.text.side_effect = lambda s="": texts.append(s) + yield bar + + with patch( + "sift_client._internal.low_level_wrappers.data.alive_bar", + _fake_alive_bar, + ): + yield texts + + +class TestGetChannelDataProgress: + """The progress bar's ``total`` reflects the number of dispatched fetches.""" + + @pytest.mark.asyncio + async def test_bar_total_equals_dispatched_task_count(self) -> None: + """One task per uncached channel → bar total matches the channel count.""" + client = DataLowLevelClient(MagicMock()) + with _capture_bar_total() as totals: + with _fake_grpc(client, {"c1": [_frame("c1")], "c2": [_frame("c2", offset=100)]}): + await client.get_channel_data( + channels=[_channel("c1"), _channel("c2")], + start_time=_NOW, + end_time=_WINDOW_END, + ignore_cache=True, + show_progress=True, + ) + assert totals == [2] + + @pytest.mark.asyncio + async def test_fully_cached_call_dispatches_no_tasks(self, tmp_path) -> None: + """A cache hit fetches nothing, so the bar total is zero.""" + client = _client_with_cache(tmp_path) + try: + with _fake_grpc(client, {"c1": [_frame("c1")]}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + with _capture_bar_total() as totals: + with _fake_grpc(client, {"c1": []}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + show_progress=True, + ) + assert totals == [0] + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_progress_text_reports_points_loaded(self) -> None: + """The situational text accumulates the point count across pages.""" + client = DataLowLevelClient(MagicMock()) + pages = [ + _frame("c1", rows=10, start=_NOW), + _frame("c1", rows=10, start=_NOW + timedelta(seconds=1)), + ] + with _capture_bar_text() as texts: + with _fake_grpc(client, {"c1": pages}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ignore_cache=True, + show_progress=True, + ) + assert any("10 pts" in t for t in texts), texts # live line after the first page + assert any("20 pts" in t for t in texts), texts # live line after both pages + # Final receipt is a past-tense summary with an average rate. + assert texts[-1].startswith("Loaded 20 points"), texts + assert "/s" in texts[-1], texts + + @pytest.mark.asyncio + async def test_progress_text_splits_fetched_and_cached(self, tmp_path) -> None: + """A mixed call reports the fetched/cached split and a fetched-only rate.""" + client = _client_with_cache(tmp_path) + try: + # Warm the cache for c1 (outside the capture block). + with _fake_grpc(client, {"c1": [_frame("c1", rows=5)]}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + # c1 is now cached; c2 is fresh, so the call mixes both. + with _capture_bar_text() as texts: + with _fake_grpc(client, {"c1": [], "c2": [_frame("c2", rows=5, offset=100)]}): + await client.get_channel_data( + channels=[_channel("c1"), _channel("c2")], + start_time=_NOW, + end_time=_WINDOW_END, + show_progress=True, + ) + final = texts[-1] + assert final.startswith("Loaded 10 points"), texts # total spans both sources + assert "5 fetched" in final, texts + assert "5 cached" in final, texts + assert "/s" in final, texts # a fetch happened, so an average rate shows + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_progress_text_all_cached_shows_no_rate(self, tmp_path) -> None: + """A fully-cached call labels the total 'all cached' and omits the rate.""" + client = _client_with_cache(tmp_path) + try: + with _fake_grpc(client, {"c1": [_frame("c1", rows=5)]}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + with _capture_bar_text() as texts: + with _fake_grpc(client, {"c1": []}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + show_progress=True, + ) + assert texts, "fully-cached call should still render a summary line" + final = texts[-1] + assert final == "Loaded 5 points from cache", texts + assert "/s" not in final, texts # nothing fetched, so no rate + finally: + client.channel_cache.store.close() diff --git a/python/lib/sift_client/resources/channels.py b/python/lib/sift_client/resources/channels.py index df5d218a9..f5cbc09d7 100644 --- a/python/lib/sift_client/resources/channels.py +++ b/python/lib/sift_client/resources/channels.py @@ -266,6 +266,7 @@ async def get_data( end_time: datetime | None = None, limit: int | None = None, ignore_cache: bool = False, + show_progress: bool | None = None, ) -> dict[str, pd.DataFrame]: """Get data for one or more channels. @@ -276,12 +277,18 @@ async def get_data( end_time: The end time to get data for. limit: The maximum number of data points to return. Will be in increments of page_size or default page size defined by the call if no page_size is provided. ignore_cache: Whether to ignore cached data and fetch fresh data from the server. + show_progress: If True, display a progress bar naming each channel as + its data is fetched. Defaults to True for sync, False for async. + Use ``sift_client.config.show_progress = False`` to disable globally. Returns: A dictionary mapping channel names to pandas DataFrames containing the channel data. """ self._ensure_data_low_level_client() + if show_progress is None: + show_progress = self._show_progress() + run_id = run._id_or_error if isinstance(run, Run) else run return await self._data_low_level_client.get_channel_data( # type: ignore channels=channels, @@ -290,6 +297,7 @@ async def get_data( end_time=end_time, max_results=limit, ignore_cache=ignore_cache, + show_progress=show_progress, ) async def get_data_as_arrow( @@ -301,6 +309,7 @@ async def get_data_as_arrow( end_time: datetime | None = None, limit: int | None = None, ignore_cache: bool = False, + show_progress: bool | None = None, ) -> dict[str, pa.Table]: """Get data for one or more channels as pyarrow tables.""" from pyarrow import Table as ArrowTable @@ -313,5 +322,6 @@ async def get_data_as_arrow( end_time=end_time, limit=limit, ignore_cache=ignore_cache, + show_progress=show_progress, ) return {k: ArrowTable.from_pandas(v) for k, v in data.items()} diff --git a/python/lib/sift_client/resources/jobs.py b/python/lib/sift_client/resources/jobs.py index 2f9387427..c75a5c3bc 100644 --- a/python/lib/sift_client/resources/jobs.py +++ b/python/lib/sift_client/resources/jobs.py @@ -299,4 +299,4 @@ async def wait_and_download( if not extract or not zipfile.is_zipfile(download_path): return [download_path] - return extract_zip(download_path, output_dir) + return extract_zip(download_path, output_dir, show_progress=show_progress) diff --git a/python/lib/sift_client/resources/sync_stubs/__init__.pyi b/python/lib/sift_client/resources/sync_stubs/__init__.pyi index d977109be..2449313dd 100644 --- a/python/lib/sift_client/resources/sync_stubs/__init__.pyi +++ b/python/lib/sift_client/resources/sync_stubs/__init__.pyi @@ -489,6 +489,7 @@ class ChannelsAPI: end_time: datetime | None = None, limit: int | None = None, ignore_cache: bool = False, + show_progress: bool | None = None, ) -> dict[str, pd.DataFrame]: """Get data for one or more channels. @@ -499,6 +500,9 @@ class ChannelsAPI: end_time: The end time to get data for. limit: The maximum number of data points to return. Will be in increments of page_size or default page size defined by the call if no page_size is provided. ignore_cache: Whether to ignore cached data and fetch fresh data from the server. + show_progress: If True, display a progress bar naming each channel as + its data is fetched. Defaults to True for sync, False for async. + Use ``sift_client.config.show_progress = False`` to disable globally. Returns: A dictionary mapping channel names to pandas DataFrames containing the channel data. @@ -514,6 +518,7 @@ class ChannelsAPI: end_time: datetime | None = None, limit: int | None = None, ignore_cache: bool = False, + show_progress: bool | None = None, ) -> dict[str, pa.Table]: """Get data for one or more channels as pyarrow tables.""" ...