The `rich.progress` support in marimo notebooks

[williambdean]

Coming from the marimo repo. There is likely a change required to support this functionality. Would like to hear thoughts on this

The rich.progress module provides progress bars for stdout. However, the result is not displayed until after the progress bar is complete which defeats the point.

I've tried many of the examples in the documentation here and I have not had luck with all I have tried.

I tried on the current main branch of the marimo (0.9.17) and version of rich 13.7.0

Example code

import marimo

__generated_with = "0.9.17"
app = marimo.App(width="full")


@app.cell
def __():
    import time
    from rich.progress import track, Progress, BarColumn, TextColumn
    from rich.table import Column

    for i in track(range(10), description="Processing..."):
        time.sleep(1)  # Simulate work being done
    return BarColumn, Column, Progress, TextColumn, i, time, track


@app.cell
def __(Progress, time):
    total = 10

    with Progress() as progress:
        task1 = progress.add_task("[red]Downloading...", total=total)
        task2 = progress.add_task("[green]Processing...", total=total)
        task3 = progress.add_task("[cyan]Cooking...", total=total)

        while not progress.finished:
            progress.update(task1, advance=0.5)
            progress.update(task2, advance=0.3)
            progress.update(task3, advance=0.9)
            time.sleep(0.02)
    return progress, task1, task2, task3, total


@app.cell
def __(Progress, progress):
    def do_work(task):
        print(task)

    with Progress(transient=True) as progress2:
        task = progress.add_task("Working", total=100)
        do_work(task)
    return do_work, progress2, task


@app.cell
def __(BarColumn, Column, Progress, TextColumn, progress, time):
    text_column = TextColumn(
        "{task.description}",
        table_column=Column(ratio=1),
    )
    bar_column = BarColumn(bar_width=None, table_column=Column(ratio=2))

    with Progress(text_column, bar_column, expand=True):
        for n in progress.track(range(10)):
            progress.print(n)
            time.sleep(0.1)
    return bar_column, n, text_column


if __name__ == "__main__":
    app.run()

Originally posted by @williambdean in https://github.com/marimo-team/marimo/issues/2846

[github-actions[bot]]

Thank you for your issue. Give us a little time to review it.

PS. You might want to check the FAQ if you haven't done so already.

This is an automated reply, generated by FAQtory

[williambdean]

Related discussion: https://github.com/Textualize/rich/discussions/3559

[willmcgugan]

This is really an issue for Marimo. I'm not even sure if you should expect it to work, as Rich is a terminal library, and I don't believe Marimo claims to be terminal compatible?

[willmcgugan]

Closing for now. Feel free to re-open if you think there is something to be done from the Rich side.

[github-actions[bot]]

I hope I helped!

Consider sponsoring my work on Rich. I give tech support for free, in addition to maintaining Rich and Textual.

If you like using Rich, you might also enjoy Textual.

Will McGugan

[williambdean]

It seems there would have to be a similar shim as this module:

https://github.com/Textualize/rich/blob/master/rich/jupyter.py

What aspects of rich does that module support?

The marimo team has a similar protocol as jupyter for displaying: https://docs.marimo.io/guides/integrating_with_marimo/displaying_objects/#option-1-implement-a-display-method

[mscolnick]

Would the rich library be open to supporting marimo natively, as it does with Jupyter?

[willmcgugan]

Would the rich library be open to supporting marimo natively, as it does with Jupyter?

Open to that.

[koaning]

@willmcgugan would you mind if I start exploring this a bit? I'm down to work on it.

[willmcgugan]

Be my guest. Let me know if you have any questions...

[koaning]

I have some first progress after patching the display function in jupyter.py. It now looks like this:

def display(segments: Iterable[Segment], text: str) -> None:
    """Render segments to Jupyter."""
    html = _render_segments(segments)
    jupyter_renderable = JupyterRenderable(html, text)
    try:
        import marimo as mo
        if mo.running_in_notebook():
            mo.output.append(jupyter_renderable)
            return
    except ModuleNotFoundError:
        pass
    try:
        from IPython.display import display as ipython_display

        ipython_display(jupyter_renderable)
    except ModuleNotFoundError:
        # Handle the case where the Console has force_jupyter=True,
        # but IPython is not installed.
        pass

I'm not sure if this approach might break, still need to look for edge cases, but ... just to check the preference @willmcgugan, I could add a marimo.py file but that feels like it might add a lot of boilerplate and overlapping code. Instead, I could also patch the jupyter code so that we do two checks, one for marimo and one for jupyter.

I'm leaning towards patching the jupyter file, but will gladly hear your opinion.

[koaning]

Ah. I think I just hit a roadblock.

https://github.com/Textualize/rich/blob/4d6d631a3d2deddf8405522d4b8c976a6d35726c/rich/live.py#L141-L143

For the terminal and Jupyter this is fine, but for marimo it's tricky because mo.output.replace() doesn't work from background threads. I did notice that manual refreshes work with my current approach, something like this:

from rich.live import Live
from rich.table import Table

table = Table()
table.add_column("Row ID")
table.add_column("Description")
table.add_column("Level")

with Live(table, refresh_per_second=4) as live:  # update 4 times a second to feel fluid
    for row in range(12):
        time.sleep(0.4)  # arbitrary delay
        # update the renderable internally
        table.add_row(f"{row}", f"description {row}", "[red]ERROR")
        live.refresh()

But without the live.refresh() here my approach will only show the first render and the last one of the Live-animation. This not only affects the Live class but also everything that depends on it, including things like the progress bar.

Quick demo for later reference that this does not work, you can confirm there's no update in marimo when you do:

from threading import Thread

def test_background():
    def update_from_thread():
        for i in range(5):
            time.sleep(0.5)
            mo.output.replace(mo.Html(f"<p>Background thread update {i}</p>"))
    
    thread = Thread(target=update_from_thread, daemon=True)
    thread.start()
    thread.join()

test_background()

It's a bummer, but it's also a road-block for the short term.

[willmcgugan]

I think patching jupyter.py is a reasonable approach. But as long as the logic remains quite isolated, I wouldn't have any issues.

Could the threading limitation of mo.output.replace be removed?

BTW, you might want to open a new issue or discussion to track this.

[koaning]

BTW, you might want to open a new issue or discussion to track this.

Fair. We have an issue internally for the mo.output.replace. Once we get round to that then I'll open up a new one.