StateViewer API#

Run this demo live in your browser ▶

StateViewer(search_func, view_func) runs search_func once and records every line it executes, including recursive calls. At each step it calls view_func with the function's local variables, then plays the whole run back:

• left: whatever view_func drew for that step;
• right: the source with the line about to run highlighted, and a table of the function's variables. Rows stay fixed, values that just changed flash, and variables that aren't set yet show not assigned;
• bottom: play/pause/step buttons, a speed slider, a badge with the current timeline value, and a timeline of the whole run. Click or drag the timeline to scrub, or focus the widget and use ← / → (Shift jumps 10), Home / End and Space.

It's made for search problems, where the interesting part is how the state changes on the way to an answer.

Quick start#

Backtracking search for N-Queens, with a view that prints the board:

from wigglystuff import StateViewer

def solve(board, row, n):
    if row == n:
        return True
    for col in range(n):
        if all(c != col and abs(c - col) != row - r for r, c in enumerate(board)):
            board.append(col)
            if solve(board, row + 1, n):
                return True
            board.pop()
    return False

def view(board, n):
    rows = ["." * c + "Q" + "." * (n - c - 1) for c in board]
    return "<pre>" + "\n".join(rows) + "</pre>"

StateViewer(solve, view, args=([], 0, 8))

args and kwargs are passed to search_func. The run is recorded when the widget is created; playing it back doesn't run your code again.

How view_func gets its values#

view_func's parameter names pick the locals it receives. def view(board, n) gets the live board and n from the frame that's running. In a recursive search that's the innermost call.

• A parameter with a default is optional: it gets the default while that local doesn't exist yet. def view(board, n, row=None, col=None) works on every step.
• While a required parameter isn't assigned yet, the previous view stays up (and the panel is empty before the first one). So give a default to anything that's assigned partway through the function, or the opening steps show nothing.
• A **kwargs parameter receives every local.

Each view is drawn the moment its step is recorded, so in-place changes like board.append(col) show up correctly.

view_func can return anything marimo can display: an HTML or SVG string, mo.Html/mo.md, a dataframe, even a matplotlib figure. Plain strings are used as HTML directly. Every distinct view is stored in the widget (repeats are stored once), so small SVG strings are the best choice for long runs.

Searches without recursion: timeline_func#

By default the timeline plots call depth: how many calls to search_func haven't returned yet. That shows a backtracking search's dives and retreats, but a loop-based search never calls itself, so its depth stays at 1 and the timeline is a flat bar.

A loop-based search keeps its progress in a data structure instead, so tell the timeline what to plot. timeline_func picks its locals by parameter name, just like view_func, and returns a number:

def search(maze, start, goal, breadth_first):
    frontier = [start]
    came_from = {start: None}
    while frontier:
        node = frontier.pop(0) if breadth_first else frontier.pop()
        if node == goal:
            break
        r, c = node
        for nxt in [(r - 1, c), (r, c + 1), (r + 1, c), (r, c - 1)]:
            if maze[nxt[0]][nxt[1]] != "#" and nxt not in came_from:
                came_from[nxt] = node
                frontier.append(nxt)
    ...

StateViewer(
    search,
    maze_view,
    args=(MAZE, (1, 1), (8, 15), True),
    timeline_func=lambda frontier: len(frontier),
    timeline_label="frontier size",
)

The badge and the caption show timeline_label. It defaults to the function's name, or "value" for a lambda. Before the parameters exist the value is 0, and it must be a finite number. In the demo, the same function plotted this way shows breadth-first search keeping a narrow band while depth-first search's stack piles up.

Driving other widgets#

view_func is optional. Without it the widget shows just the code, variables and timeline, and you draw the state with any other widget instead:

• step is synced back to Python while you play and scrub, at most once every sync_throttle_ms (default 100; 0 syncs every step).
• viewer.state holds the real local values at that step, not reprs. They're copied when the step is recorded, so later in-place changes don't affect them.
• variables=["board", "row"] limits the table to those names in that order. state still holds every local.

viewer = StateViewer(solve, args=([], 0, 8), variables=["board", "row", "col"])
viewer.observe(lambda _: redraw(viewer.state["board"]), names="step")

In marimo, wrap the viewer in mo.ui.anywidget and update another widget from a cell that reads viewer.widget.state. Assign to that widget's traitlets rather than creating a new one, so it keeps its own state. The demo uses this to drive a rotatable three.js board.

Reading the trace from Python#

• viewer.steps is the full trace, one record per step:
  • line: the line about to run (1 = the def line; decorators are left out of code);
  • depth: call depth;
  • event: "line", "return" or "raise";
  • locals: reprs of the shown variables;
  • view: an index into views (-1 = nothing drawn yet);
  • value: the timeline_func value, if you passed one;
  • returned: on return steps, the returned value's repr.
• viewer.step, viewer.current (the record at step), viewer.state and viewer.n_steps.
• viewer.returned is the repr of the result, viewer.truncated says whether max_steps was reached, and viewer.error is set if search_func, view_func or timeline_func raised. Errors are shown in the widget rather than raised, with the line they happened on.

Limits#

• Only search_func itself is recorded. Calls into helper functions run but aren't recorded, and neither are functions defined inside it or the generator expression in an all(...) check. If the interesting work happens in a helper, trace the helper.
• search_func needs Python source you can inspect: a plain def in a file or notebook cell works best. Builtins and bound methods (obj.solve) raise a TypeError. For a method, pass Class.solve and the object in args.
• Generator functions and async def functions record zero steps: calling them only creates the generator or coroutine and doesn't run the body.
• max_steps (default 20,000) stops recording a runaway search and sets truncated. only_changes=True keeps only steps whose view (or, without a view, whose shown values) differs from the previous step.
• Threads started by search_func aren't recorded; only the thread that creates the widget is.
• Every step copies the locals (for state) and repr's them (for the table), so very large objects make recording slower.

See also: LiveEdit for a table of every loop iteration of one run, and AsyncFlow for async tasks.

from wigglystuff import StateViewer

def solve(board, row, n):
    if row == n:
        return True
    for col in range(n):
        if all(c != col and abs(c - col) != row - r for r, c in enumerate(board)):
            board.append(col)
            if solve(board, row + 1, n):
                return True
            board.pop()
    return False

def view(board, n):
    return "<pre>" + "\n".join("." * c + "Q" + "." * (n - c - 1) for c in board) + "</pre>"

StateViewer(solve, view, args=([], 0, 6))

viewer = StateViewer(solve, args=([], 0, 6), variables=["board", "row"])
viewer.observe(lambda _: print(viewer.state["board"]), names="step")

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428

def __init__(
    self,
    search_func: Callable[..., Any],
    view_func: Callable[..., Any] | None = None,
    *,
    args: tuple[Any, ...] = (),
    kwargs: dict[str, Any] | None = None,
    max_steps: int = 20_000,
    only_changes: bool = False,
    interval_ms: int = 80,
    loop: bool = False,
    height: int = 0,
    variables: list[str] | None = None,
    sync_throttle_ms: int = 100,
    timeline_func: Callable[..., Any] | None = None,
    timeline_label: str | None = None,
    **widget_kwargs: Any,
) -> None:
    """Create a StateViewer.

    Args:
        search_func: The function to trace. Needs inspectable Python source.
        view_func: Draws the state; its parameter names pick the locals it gets.
            Optional: leave it out to only show code, variables and timeline.
        args: Positional arguments for ``search_func``.
        kwargs: Keyword arguments for ``search_func``.
        max_steps: Stop tracing after this many steps (sets ``truncated``).
        only_changes: Drop steps whose view (or, without a view, whose shown
            values) is identical to the previous step's.
        interval_ms: Milliseconds per step while playing.
        loop: Restart from the first step when playback reaches the end.
        height: Height cap in pixels (0 = fit content).
        variables: Local names to show in the variables table, in this order.
            Defaults to every local. ``state`` always holds all of them.
        sync_throttle_ms: Cap on how often ``step`` syncs to Python while
            playing or scrubbing, in milliseconds (0 = every step).
        timeline_func: Returns the number the timeline plots at each step;
            its parameter names pick the locals, like ``view_func``.
            Defaults to the call depth of ``search_func``, which stays flat
            for a search without recursion, so pass e.g.
            ``lambda queue: len(queue)`` there.
        timeline_label: Name shown for the timeline value. Defaults to
            ``timeline_func``'s name, or "value" for a lambda.
        **widget_kwargs: Forwarded to ``anywidget.AnyWidget``.
    """
    if max_steps <= 0:
        raise ValueError("max_steps must be positive.")
    if interval_ms <= 0:
        raise ValueError("interval_ms must be positive.")
    if view_func is not None and not callable(view_func):
        raise TypeError("view_func must be callable.")
    if timeline_func is not None and not callable(timeline_func):
        raise TypeError("timeline_func must be callable.")
    if timeline_label is None:
        name = getattr(timeline_func, "__name__", "<lambda>")
        if timeline_func is None:
            timeline_label = "call depth"
        else:
            timeline_label = "value" if name == "<lambda>" else name
    code, lines, skipped = _source_lines(search_func)
    fn_code = search_func.__code__
    # dict.fromkeys dedupes: a local captured by a closure or generator
    # expression is listed in both co_varnames and co_cellvars.
    all_names = list(dict.fromkeys((*fn_code.co_varnames, *fn_code.co_cellvars)))
    if variables is not None:
        unknown = [name for name in variables if name not in all_names]
        if unknown:
            raise ValueError(
                f"Unknown variables {unknown}; {search_func.__name__} has {all_names}."
            )
    recorder = _Recorder(
        search_func, view_func, max_steps, skipped, variables, timeline_func
    )
    returned = None
    error = None
    previous_trace = sys.gettrace()
    sys.settrace(recorder.global_trace)
    try:
        returned = _short_repr(search_func(*args, **(kwargs or {})))
    except _StepLimit:
        pass
    except Exception as exc:  # noqa: BLE001 - search errors are data, not crashes.
        error = self._error_payload(exc, recorder, "search_func")
    finally:
        sys.settrace(previous_trace)
    if recorder.callback_error is not None:
        error = self._error_payload(
            recorder.callback_error, recorder, recorder.callback_source
        )

    steps = recorder.steps
    states = recorder.states
    if only_changes:
        key = "view" if view_func is not None else "locals"
        keep = [i for i in range(len(steps)) if i == 0 or steps[i][key] != steps[i - 1][key]]
        steps = [steps[i] for i in keep]
        states = [states[i] for i in keep]
    self._states = states
    if variables is None:
        seen = {name for s in steps for name in s["locals"]}
        variables = [name for name in all_names if name in seen]
    super().__init__(
        code=code,
        lines=lines,
        steps=steps,
        views=recorder.views,
        variables=variables,
        interval_ms=interval_ms,
        loop=loop,
        returned=returned,
        error=error,
        truncated=recorder.truncated,
        height=height,
        sync_throttle_ms=sync_throttle_ms,
        timeline_label=timeline_label,
        **widget_kwargs,
    )

current: dict[str, Any] | None

n_steps: int

state: dict[str, Any] | None

Synced traitlets#