`rich.progress` support

[williambdean]

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()

[mscolnick]

We might need to add marimo support to rich directly https://github.com/Textualize/rich/blob/43d3b04725ab9731727fb1126e35980c62f32377/rich/console.py#L517-L534

Would you be open to create an issue on rich?

[williambdean]

We might need to add marimo support to rich directly https://github.com/Textualize/rich/blob/43d3b04725ab9731727fb1126e35980c62f32377/rich/console.py#L517-L534

Would you be open to create an issue on rich?

Sure. What is the logic to detect for marimo?

[mscolnick]

@wd60622, there is this logic from another open-source project:

https://github.com/run-house/runhouse/blob/2f80fe25879d7df0f3b8432cf62b3c0519d7ea36/runhouse/resources/module.py#L1073-L1080

        # Check if running in a marimo notebook
        try:
            import marimo
            return marimo.running_in_notebook()
        except (ImportError, ModuleNotFoundError):
            # marimo not installed
            return False

[smutch]

In the short term, I've been using this to monkey-patch Rich and make it think it is outputting to a terminal:

import rich.console


_orig_console = rich.console.Console


class Console(_orig_console):
    def __init__(self, *args, **kwargs):
        kwargs["force_terminal"] = True
        super().__init__(*args, **kwargs)


rich.console.Console = Console

This allows the progress bar to update in real time.

[williambdean]

Thanks for that lead. That doesn't work for my use-case unfortunately.

I'm doing something like this

PyMC x Marimo

import marimo

__generated_with = "0.12.8"
app = marimo.App()


@app.cell
def _():
    import pymc as pm
    import numpy as np

    seed = sum(map(ord, "Marimo x PyMC"))
    rng = np.random.default_rng(10)

    return np, pm, rng, seed


@app.cell
def _():
    nuts_sampler = "pymc"
    return (nuts_sampler,)


@app.cell
def _(model, nuts_sampler, pm):
    with model:
        idata = pm.sample(10_000, nuts_sampler=nuts_sampler)

    idata
    return (idata,)


@app.cell
def _(pm, rng):
    coords = {"date": [1, 2, 3]}
    model = pm.Model(coords=coords)

    data = rng.normal(loc=0, scale=3, size=1000)

    with model:
        mu = pm.Normal("mu")
        sigma = pm.HalfNormal("sigma")
        nu = pm.HalfNormal("nu")
        pm.StudentT("obs", nu=nu, mu=mu, sigma=sigma, observed=data)
    return coords, data, model, mu, nu, sigma


@app.cell
def _(model):
    model
    return


@app.cell
def _():
    import rich.console

    _orig_console = rich.console.Console

    class Console(_orig_console):
        def __init__(self, *args, **kwargs):
            kwargs["force_terminal"] = True
            super().__init__(*args, **kwargs)

    rich.console.Console = Console
    return Console, rich


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

[smutch]

I think this is maybe due to the ordering of things. It should work if you place the monkey patching code in the same cell as where you import pymc, and make sure it comes before the pymc import.

Having said that, I just tried it with pymc and instead of clearing and redrawing the parallel progressbars, it just prints a new copy of them for every iteration. :(

[mscolnick]

it would be great to push support directly into rich.console. would anyone be interested in either filing an issue or making the contribution to that repo?

[williambdean]

That hack doesnt work. However, I will raise an issue on the rich repo

[williambdean]

The answer over there is that it should be implemented in marimo

[mscolnick]

Given all their shims and customization for rendering in Jupyter, I would imagine we would need the same for marimo