From b91d54e45b596b78f3dc972993ff7c6467958a71 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Sat, 21 Jun 2025 20:12:25 -0400 Subject: [PATCH 1/3] trying out with mo.output --- python/nutpie/sample.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py index bbb7b9c..238a22c 100644 --- a/python/nutpie/sample.py +++ b/python/nutpie/sample.py @@ -267,6 +267,15 @@ def _trace_to_arviz(traces, n_tune, shapes, **kwargs): """ +def in_marimo_notebook() -> bool: + try: + import marimo as mo + + return mo.running_in_notebook() + except ImportError: + return False + + # Adapted from fastprogress def in_notebook(): def in_colab(): @@ -362,6 +371,27 @@ def callback(formatted): self._html = formatted self.display_id.update(self) + progress_type = _lib.ProgressType.template_callback( + progress_rate, progress_template, cores, callback + ) + elif in_marimo_notebook(): + import marimo as mo + + if progress_template is None: + progress_template = _progress_template + + if progress_style is None: + progress_style = _progress_style + + self._html = "" + + mo.output.replace(mo.Html("Sampling is about to start...")) + + def callback(formatted): + self._html = formatted + html = mo.Html(f"{progress_style}\n{formatted}") + mo.output.replace(html) + progress_type = _lib.ProgressType.template_callback( progress_rate, progress_template, cores, callback ) From baf04f760f7f10e39bf9a8ed3f1e3307cc9ef549 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Sat, 21 Jun 2025 23:22:53 -0400 Subject: [PATCH 2/3] getting context and stream before going to callback --- python/nutpie/sample.py | 46 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 44 insertions(+), 2 deletions(-) diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py index 238a22c..817a680 100644 --- a/python/nutpie/sample.py +++ b/python/nutpie/sample.py @@ -276,6 +276,47 @@ def in_marimo_notebook() -> bool: return False +def write_internal(cell_id, stream, value: object) -> None: + from marimo._output import formatting + from marimo._messaging.ops import CellOp + from marimo._messaging.tracebacks import write_traceback + from marimo._messaging.cell_output import CellChannel + + output = formatting.try_format(value) + if output.traceback is not None: + write_traceback(output.traceback) + CellOp.broadcast_output( + channel=CellChannel.OUTPUT, + mimetype=output.mimetype, + data=output.data, + cell_id=cell_id, + status=None, + stream=stream, + ) + + +def create_replace(): + from marimo._runtime.context import get_context + from marimo._runtime.context.types import ContextNotInitializedError + from marimo._output import formatting + + try: + ctx = get_context() + except ContextNotInitializedError: + return + + cell_id = ctx.execution_context.cell_id + execution_context = ctx.execution_context + stream = ctx.stream + + def replace(value): + execution_context.output = [formatting.as_html(value)] + + write_internal(cell_id=cell_id, value=value, stream=stream) + + return replace + + # Adapted from fastprogress def in_notebook(): def in_colab(): @@ -385,12 +426,13 @@ def callback(formatted): self._html = "" - mo.output.replace(mo.Html("Sampling is about to start...")) + mo.output.clear() + my_replace = create_replace() def callback(formatted): self._html = formatted html = mo.Html(f"{progress_style}\n{formatted}") - mo.output.replace(html) + my_replace(html) progress_type = _lib.ProgressType.template_callback( progress_rate, progress_template, cores, callback From 72a5b27eb8c4bde8d254bb988ad3d5d35afbdb7e Mon Sep 17 00:00:00 2001 From: Will Dean Date: Sun, 22 Jun 2025 06:49:38 -0400 Subject: [PATCH 3/3] use private functions --- python/nutpie/sample.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py index 817a680..0655173 100644 --- a/python/nutpie/sample.py +++ b/python/nutpie/sample.py @@ -276,7 +276,8 @@ def in_marimo_notebook() -> bool: return False -def write_internal(cell_id, stream, value: object) -> None: +def _mo_write_internal(cell_id, stream, value: object) -> None: + """Write to marimo cell given cell_id and stream.""" from marimo._output import formatting from marimo._messaging.ops import CellOp from marimo._messaging.tracebacks import write_traceback @@ -295,7 +296,8 @@ def write_internal(cell_id, stream, value: object) -> None: ) -def create_replace(): +def _mo_create_replace(): + """Create mo.output.replace with current context pinned.""" from marimo._runtime.context import get_context from marimo._runtime.context.types import ContextNotInitializedError from marimo._output import formatting @@ -312,7 +314,7 @@ def create_replace(): def replace(value): execution_context.output = [formatting.as_html(value)] - write_internal(cell_id=cell_id, value=value, stream=stream) + _mo_write_internal(cell_id=cell_id, value=value, stream=stream) return replace @@ -427,12 +429,12 @@ def callback(formatted): self._html = "" mo.output.clear() - my_replace = create_replace() + mo_output_replace = _mo_create_replace() def callback(formatted): self._html = formatted html = mo.Html(f"{progress_style}\n{formatted}") - my_replace(html) + mo_output_replace(html) progress_type = _lib.ProgressType.template_callback( progress_rate, progress_template, cores, callback