diff --git a/frontend/src/plugins/layout/ProgressPlugin.tsx b/frontend/src/plugins/layout/ProgressPlugin.tsx index 9b9969420dc..ce992139bd7 100644 --- a/frontend/src/plugins/layout/ProgressPlugin.tsx +++ b/frontend/src/plugins/layout/ProgressPlugin.tsx @@ -1,7 +1,7 @@ /* Copyright 2026 Marimo. All rights reserved. */ import humanizeDuration from "humanize-duration"; -import { Loader2Icon } from "lucide-react"; +import { CheckCircle2Icon, Loader2Icon } from "lucide-react"; import React, { type JSX, type PropsWithChildren } from "react"; import { z } from "zod"; import { Progress } from "@/components/ui/progress"; @@ -38,6 +38,10 @@ interface Data { * The rate of progress in items per second. */ rate?: number; + /** + * Whether the progress indicator is done (shows a checkmark). + */ + done?: boolean; } export class ProgressPlugin implements IStatelessPlugin { @@ -50,6 +54,7 @@ export class ProgressPlugin implements IStatelessPlugin { total: z.number().optional(), eta: z.number().optional(), rate: z.number().optional(), + done: z.boolean().optional(), }); render(props: IStatelessPluginProps): JSX.Element { @@ -64,11 +69,17 @@ export const ProgressComponent = ({ total, eta, rate, + done, }: PropsWithChildren): JSX.Element => { const alignment = typeof progress === "number" ? "items-start" : "items-center"; const renderProgress = () => { + // When done, show a checkmark + if (done) { + return ; + } + // With a known total, show a progress bar. if (typeof progress === "number" && total != null && total > 0) { return ( diff --git a/marimo/_plugins/stateless/status/_progress.py b/marimo/_plugins/stateless/status/_progress.py index 6311acb5a1b..684b97c961a 100644 --- a/marimo/_plugins/stateless/status/_progress.py +++ b/marimo/_plugins/stateless/status/_progress.py @@ -50,6 +50,7 @@ def __init__( self.total = total self.current = 0 self.closed = False + self._is_done = False # We show a loading spinner if total not known self.loading_spinner = total is None self.show_rate = show_rate @@ -116,6 +117,34 @@ def close(self) -> None: output.flush() # Flush one last time before closing self.closed = True + def mark_done( + self, + title: str | None = None, + subtitle: str | None = None, + ) -> None: + """Mark the progress indicator as done. Thread-safe. + + Transitions the spinner to a completed (checkmark) state. + Optionally updates title and subtitle. + + Args: + title (str, optional): New title. Defaults to None. + subtitle (str, optional): New subtitle. Defaults to None. + """ + with self._lock: + if self.closed: + raise RuntimeError( + "Progress indicators cannot be updated after exiting " + "the context manager that created them. " + ) + self._is_done = True + if title is not None: + self.title = title + if subtitle is not None: + self.subtitle = subtitle + self._text = self._get_text() + self.debounced_flush() + def _get_text(self) -> str: return build_stateless_plugin( component_name="marimo-progress", @@ -129,6 +158,7 @@ def _get_text(self) -> str: "progress": True if self.loading_spinner else self.current, "rate": self._get_rate(), "eta": self._get_eta(), + "done": True if self._is_done else None, } ), ) @@ -227,6 +257,28 @@ def update( """ super().update_progress(increment=1, title=title, subtitle=subtitle) + def done( + self, title: str | None = None, subtitle: str | None = None + ) -> None: + """Mark the spinner as done, transitioning to a checkmark state. + + After calling `done()`, the spinner shows a checkmark instead of + the loading animation. Optionally update the title and subtitle + to reflect the completed state. + + Examples: + ```python + with mo.status.spinner("Loading ...") as _spinner: + data = expensive_function() + _spinner.done(title="Done!") + ``` + + Args: + title (str, optional): New title. Defaults to None. + subtitle (str, optional): New subtitle. Defaults to None. + """ + super().mark_done(title=title, subtitle=subtitle) + @mddoc class spinner: @@ -273,9 +325,10 @@ def __enter__(self) -> Spinner: return self.spinner def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: - if self.remove_on_exit: + if self.remove_on_exit or exc_type is not None: self.spinner.clear() - # TODO(akshayka): else consider transitioning to a done state + else: + self.spinner.mark_done() self.spinner.close() def _mime_(self) -> tuple[KnownMimeType, str]: diff --git a/tests/_plugins/stateless/status/test_progress.py b/tests/_plugins/stateless/status/test_progress.py index 81240621dff..09fa1f5eaa3 100644 --- a/tests/_plugins/stateless/status/test_progress.py +++ b/tests/_plugins/stateless/status/test_progress.py @@ -8,6 +8,7 @@ import pytest from marimo._plugins.stateless.status._progress import ( + Spinner, _Progress, progress_bar, spinner, @@ -148,6 +149,71 @@ def test_spinner_without_context(): _spinner.update(subtitle="Crunching numbers ...") +@patch("marimo._runtime.output._output.flush") +def test_spinner_done(mock_flush: Any) -> None: + del mock_flush + s = Spinner(title="Loading", subtitle="Please wait") + assert s._is_done is False + + s.done(title="All done", subtitle="Completed!") + assert s._is_done is True + assert s.title == "All done" + assert s.subtitle == "Completed!" + + +@patch("marimo._runtime.output._output.flush") +def test_spinner_done_default_title(mock_flush: Any) -> None: + del mock_flush + s = Spinner(title="Loading", subtitle="Please wait") + s.done() + assert s._is_done is True + assert s.title == "Loading" + assert s.subtitle == "Please wait" + + +@patch("marimo._runtime.output._output.flush") +def test_spinner_done_closed_raises(mock_flush: Any) -> None: + del mock_flush + s = Spinner(title="Loading", subtitle="Please wait") + s.close() + with pytest.raises(RuntimeError): + s.done() + + +def test_spinner_context_manager_remove_on_exit(): + assert runtime_context_installed() is False + + with spinner("Test", remove_on_exit=True) as _spinner: + _spinner.update(subtitle="Working...") + # After exit with remove_on_exit=True, spinner is cleared + assert _spinner.closed is True + + +def test_spinner_context_manager_keep_on_exit(): + assert runtime_context_installed() is False + + with spinner("Test", remove_on_exit=False) as _spinner: + _spinner.update(subtitle="Working...") + # After exit with remove_on_exit=False, spinner transitions to done state + assert _spinner._is_done is True + assert _spinner.closed is True + + +@patch("marimo._runtime.output._output.remove") +def test_spinner_context_manager_exception_does_not_mark_done( + mock_remove: Any, +) -> None: + assert runtime_context_installed() is False + + with pytest.raises(ValueError, match="Failed"): + with spinner("Test", remove_on_exit=False) as _spinner: + raise ValueError("Failed") + + assert _spinner._is_done is False + assert _spinner.closed is True + mock_remove.assert_called_once_with(_spinner) + + def test_progress_without_context(): assert runtime_context_installed() is False