add layout and fix rendering using python script, not notebook env

docs/myst.yml

@@ -18,6 +18,7 @@ project:
           children:
             - file: getting-started/quickstart/candlestick.ipynb
             - file: getting-started/quickstart/dataframe.ipynb
+    - file: layout.ipynb
     - file: customization.ipynb
     - file: rendering.ipynb
     - file: cli_commands.md

pfund_plot/__init__.py

@@ -13,7 +13,7 @@
     candlestick_plot as ohlc,
     candlestick_plot as kline,
 )
-from pfund_plot.layout import layout_plot as layout
+from pfund_plot.layout import layout
 
 
 hvplot.extension('bokeh', 'plotly')

pfund_plot/layout.py

@@ -1,17 +1,74 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 if TYPE_CHECKING:
     from pfund_plot.types.core import tFigure
-    from pfund_plot.types.literals import tDisplayMode
 
-import panel as pn
-
+from contextlib import contextmanager
+
+from panel.layout.gridstack import GridStack
+
+from pfund_plot.state import state
 from pfund_plot.renderer import render
 
 
-def layout_plot(
+__all__ = ['layout']
+
+
+@contextmanager
+def layout(
+    streaming: bool = False,
+    display_mode: Literal['browser', 'desktop'] = 'browser', 
+    num_cols: int = 3,
+    allow_drag: bool = True,
+    # FIXME: plots can't be resized when using GridStack, not sure if it's a bottleneck or a bug
+    # allow_resize: bool = True,
+):
+    assert display_mode.lower() in ['browser', 'desktop'], "display_mode must be 'browser' or 'desktop'"
+
+    # Setup state
+    state.layout.in_layout = True
+    state.layout.streaming = streaming
+    state.layout.components = []
+
+    try:
+        yield state.layout
+    finally:
+        components = state.layout.components
+        return _layout_plot(
+            *components,
+            display_mode=display_mode,
+            num_cols=num_cols,
+            allow_drag=allow_drag,
+            # allow_resize=allow_resize,
+        )
+
+
+def _layout_plot(
     *figs: tFigure,
-    display_mode: tDisplayMode = 'notebook',
+    display_mode: Literal['browser', 'desktop'] = 'browser',
+    num_cols: int = 3,
+    allow_drag: bool = True,
+    # allow_resize: bool = False,
 ):
-    # use gridstack to layout the plots automatically
-    return render(combined_fig, display_mode)
+    if not state.layout.in_layout:
+        raise ValueError("layout_plot() must be called within a layout context manager")
+
+    gstack = GridStack(
+        sizing_mode='stretch_both',
+        allow_drag=allow_drag,
+        allow_resize=False,
+    )
+
+    # standardize the figures
+    max_height = max(fig.height or 0 for fig in figs)
+    for fig in figs:
+        fig.param.update(height=max_height, width=None, sizing_mode='stretch_width')
+
+    num_figs = len(figs)
+    # num_rows = math.ceil(num_figs / num_cols)
+    for i in range(num_figs):
+        gstack[i // num_cols, i % num_cols] = figs[i]
+
+    periodic_callbacks = state.layout.periodic_callbacks
+    state.reset_layout()
+    return render(gstack, display_mode, periodic_callbacks=periodic_callbacks)

pfund_plot/plots/candlestick.py

@@ -17,6 +17,7 @@
 from pfund_plot.utils.validate import validate_data_type
 from pfund_plot.utils.utils import get_sizing_mode
 from pfund_plot.renderer import render
+from pfund_plot.state import state
 
 
 __all__ = ['candlestick_plot']
@@ -130,6 +131,8 @@ def candlestick_plot(
 
 
     display_mode = DisplayMode[display_mode.lower()]
+    if state.layout.in_layout:
+        streaming = streaming or state.layout.streaming
     data_type: DataType = validate_data_type(data, streaming, import_hvplot=True)
     if data_type == DataType.datafeed:
         # TODO: get streaming data in the format of dataframe, and then call _validate_df
@@ -236,4 +239,4 @@ def _update_plot():
             height=height,
             width=width,
         )
-        return render(fig, display_mode, periodic_callback=periodic_callback)
+        return render(fig, display_mode, periodic_callbacks=[periodic_callback])

pfund_plot/plots/dataframe.py

@@ -17,6 +17,7 @@
 from pfund_plot.utils.validate import validate_data_type
 from pfund_plot.utils.utils import get_notebook_type, get_sizing_mode
 from pfund_plot.renderer import render
+from pfund_plot.state import state
 
 
 __all__ = ['dataframe_plot']
@@ -71,6 +72,8 @@ def dataframe_plot(
     '''
 
     display_mode, backend = DisplayMode[display_mode.lower()], DataFrameBackend[backend.lower()]
+    if state.layout.in_layout:
+        streaming = streaming or state.layout.streaming
     data_type: DataType = validate_data_type(data, streaming, import_hvplot=False)
     if data_type == DataType.datafeed:
         # TODO: get streaming data in the format of dataframe, and then call _validate_df
@@ -79,9 +82,10 @@ def dataframe_plot(
     else:
         df = data
     df = convert_to_pandas_df(df)
-    use_iframe_in_notebook = (backend == DataFrameBackend.perspective)
-    iframe_style = None
+
+    use_iframe_in_notebook, iframe_style = False, None
     if display_mode == DisplayMode.notebook:
+        use_iframe_in_notebook = (backend == DataFrameBackend.perspective)
         height = height or DEFAULT_HEIGHT_FOR_NOTEBOOK
     if 'sizing_mode' not in kwargs:
         kwargs['sizing_mode'] = get_sizing_mode(height, width)
@@ -96,15 +100,17 @@ def dataframe_plot(
             )
             max_streaming_data = SUGGESTED_MIN_STREAMING_DATA_FOR_TABULATOR
         notebook_type: NotebookType = get_notebook_type()
+        # FIXME: this is a workaround for a bug in panel Tabulator, see if panel will fix it, or create a github issue
+        if display_mode == DisplayMode.notebook and notebook_type in [NotebookType.jupyter, NotebookType.marimo]:
+            pagination = 'remote'
+        else:
+            pagination = 'local'
         table: Widget = pn.widgets.Tabulator(
             df,
             page_size=page_size if not max_streaming_data else max(page_size, max_streaming_data), 
             header_filters=header_filters,
             disabled=True,  # not allow user to edit the table
-            # HACK: jupyter notebook is running in a server, use remote pagination to work around the update error when streaming=True
-            # the error is: "ValueError: Must have equal len keys and value when setting with an iterable"
-            # FIXME: this is a workaround for a bug in panel Tabulator, see if panel will fix it, or create a github issue
-            pagination='local' if notebook_type == NotebookType.vscode else 'remote',
+            pagination=pagination,
             formatters={
                 # NOTE: %f somehow doesn't work for microseconds, and %N (nanoseconds) only preserves up to milliseconds precision
                 # so just use %3N to display milliseconds precision
@@ -148,7 +154,8 @@ def _update_table():
             # TEMP: fake streaming data
             # NOTE: must be pandas dataframe, pandas series, or dict
             new_data = df.tail(1)
-            new_data['symbol'] = f'AAPL_{n}'
+            if 'symbol' in new_data.columns:
+                new_data['symbol'] = f'AAPL_{n}'
             n += 1
 
             if backend == DataFrameBackend.tabulator:
@@ -160,7 +167,7 @@ def _update_table():
     return render(
         table, 
         display_mode, 
-        periodic_callback=periodic_callback, 
+        periodic_callbacks=[periodic_callback], 
         use_iframe_in_notebook=use_iframe_in_notebook,
         iframe_style=iframe_style,
     )

pfund_plot/renderer.py

@@ -5,18 +5,19 @@
     from panel.widgets import Widget
     from panel.pane import Pane
     from panel.io.threads import StoppableThread
-    from panel.io.state import PeriodicCallback
     from pfund_plot.types.core import tOutput
 
 import time
 from threading import Thread
 from multiprocessing import Process, Event
 
 import panel as pn
+from panel.io.callbacks import PeriodicCallback
 
 from pfund import print_warning
 from pfund_plot.const.enums import DisplayMode, NotebookType
 from pfund_plot.utils.utils import get_notebook_type, get_free_port
+from pfund_plot.state import state
 
 
 def run_webview(title: str, port: int, window_ready: Event):
@@ -31,32 +32,41 @@ def run_webview(title: str, port: int, window_ready: Event):
     wv.start()
 
 
-def _handle_periodic_callback(periodic_callback: PeriodicCallback | None):
+def _handle_periodic_callback(periodic_callback: PeriodicCallback | None, notebook_type: NotebookType | None):
     # the main idea is don't use the thread created by periodic_callback.start(), instead create a marimo thread to stream updates
     def _handle_marimo_streaming(periodic_callback: PeriodicCallback):
         import marimo as mo
         get_streaming_active, set_streaming_active = mo.state(True)
 
         def stream_updates():
+            time.sleep(1)  # HACK: wait some time to avoid data loss
             while get_streaming_active():  # Use the getter function
                 periodic_callback.callback()
                 time.sleep(periodic_callback.period / 1000)
 
         stream_thread = mo.Thread(target=stream_updates, daemon=True)
         stream_thread.start()
-
-    notebook_type: NotebookType = get_notebook_type()
+
     if periodic_callback:
-        if notebook_type == NotebookType.marimo:
-            _handle_marimo_streaming(periodic_callback)
+        # postpone the periodic callback until the layout plot is called; otherwise there will be data loss
+        if state.layout.in_layout:
+            state.layout.add_periodic_callback(periodic_callback)
         else:
-            periodic_callback.start()
+            if notebook_type == NotebookType.marimo:
+                _handle_marimo_streaming(periodic_callback)
+            else:
+                periodic_callback.start()
+
+
+def run_callbacks(periodic_callbacks: list[PeriodicCallback], notebook_type: NotebookType | None):
+    for callback in periodic_callbacks:
+        _handle_periodic_callback(callback, notebook_type)
 
 
 def render(
     fig: Panel | Pane | Widget,
     display_mode: Literal["notebook", "browser", "desktop"] | DisplayMode,
-    periodic_callback: PeriodicCallback | None = None,
+    periodic_callbacks: list[PeriodicCallback] | PeriodicCallback | None = None,
     use_iframe_in_notebook: bool = False,
     iframe_style: str | None = None,
 ) -> tOutput:
@@ -75,16 +85,25 @@ def render(
         iframe_style: the style of the iframe when use_iframe_in_notebook is True.
     '''
     if isinstance(display_mode, str):
-        display_mode = DisplayMode[display_mode.lower()]
-
+        display_mode = DisplayMode[display_mode.lower()]    
+
+    if isinstance(periodic_callbacks, PeriodicCallback):
+        periodic_callbacks = [periodic_callbacks]
+
+    notebook_type: NotebookType | None = get_notebook_type()
+    # NOTE: handling differs between notebook environment and python script
+    is_notebook_env = (notebook_type is not None)
+
     if display_mode == DisplayMode.notebook:
         if not use_iframe_in_notebook:
-            panel_fig: Panel | Widget = fig
+            panel_fig: Panel | Pane | Widget = fig
+            run_callbacks(periodic_callbacks, notebook_type)
         else:
             if iframe_style is None:
                 print_warning("No iframe_style is provided for iframe in notebook")
             port = get_free_port()
-            server: StoppableThread = pn.serve(fig, show=False, threaded=True, port=port)
+            if notebook_type == NotebookType.jupyter:
+                print_warning(f"If the plot can't be displayed, try to use 'from IPython.display import IFrame' and 'IFrame(src='http://localhost:{port}', width=..., height=...)'")
             panel_fig: Pane = pn.pane.HTML(
                 f'''
                 <iframe 
@@ -93,41 +112,60 @@ def render(
                 </iframe>
                 ''',
             )
-            # let iframe inherit the height, width and sizing_mode from the figure
+            # let pane HTML inherit the height, width and sizing_mode from the figure, useful in layout_plot()
             panel_fig.param.update(
                 height=fig.height,
                 width=fig.width,
                 sizing_mode=fig.sizing_mode,
             )
-        _handle_periodic_callback(periodic_callback)
+            if is_notebook_env:
+                server: StoppableThread = pn.serve(fig, show=False, threaded=True, port=port)
+                run_callbacks(periodic_callbacks, notebook_type)
+            else:  # only happens when running layout_plot() where components are all using display_mode="notebook" in a python script 
+                def run_server():
+                    run_callbacks(periodic_callbacks, notebook_type)
+                    pn.serve(fig, show=False, threaded=True, port=port)  # this will block the main thread
+                thread = Thread(target=run_server, daemon=True)
+                thread.start()
         return panel_fig
     elif display_mode == DisplayMode.browser:
-        server: StoppableThread = pn.serve(fig, show=True, threaded=True)
-        _handle_periodic_callback(periodic_callback)
-        return server
+        if is_notebook_env:
+            server: StoppableThread = pn.serve(fig, show=True, threaded=True)
+            run_callbacks(periodic_callbacks, notebook_type)
+            return server
+        else:  # run in a python script
+            run_callbacks(periodic_callbacks, notebook_type)
+            pn.serve(fig, show=True, threaded=False)  # this will block the main thread
     elif display_mode == DisplayMode.desktop:
         port = get_free_port()
-        server: StoppableThread = pn.serve(fig, show=False, threaded=True, port=port)
         title = getattr(fig, 'name', "PFund Plot")
         window_ready = Event()
-        def run_process():
-            try:
-                # NOTE: need to run in a separate process, otherwise jupyter notebook will hang after closing the webview window
-                process = Process(target=run_webview, name=title, args=(title, port, window_ready,), daemon=True)
-                process.start()
-                process.join()
-            except Exception as e:
-                print(f"An error occurred: {e}")
-            finally:
-                server.stop()
-        # NOTE: need to run the process in a separate thread, otherwise periodic callbacks when streaming=True won't work
-        # because process.join() will block the thread
-        thread = Thread(target=run_process, daemon=True)
-        thread.start()
-
-        # wait for the window to be ready before starting the periodic callback to prevent data loss when streaming=True
-        window_ready.wait()
-        _handle_periodic_callback(periodic_callback)
-        return server
+        if is_notebook_env:
+            server: StoppableThread = pn.serve(fig, show=False, threaded=True, port=port)
+            def run_process():
+                try:
+                    process = Process(target=run_webview, name=title, args=(title, port, window_ready,), daemon=True)
+                    process.start()
+                    process.join()
+                except Exception as e:
+                    print(f"An error occurred: {e}")
+                finally:
+                    server.stop()
+            thread = Thread(target=run_process, daemon=True)
+            thread.start()
+            # wait for the window to be ready before starting the periodic callback to prevent data loss when streaming=True
+            window_ready.wait()
+            run_callbacks(periodic_callbacks, notebook_type)
+            return server
+        else:
+            process = Process(target=run_webview, name=title, args=(title, port, window_ready,), daemon=True)
+            process.start()
+            def run_server():
+                run_callbacks(periodic_callbacks, notebook_type)
+                pn.serve(fig, show=False, threaded=False, port=port)  # this will block the main thread
+            window_ready.wait()
+            thread = Thread(target=run_server, daemon=True)
+            thread.start()
+            process.join()
     else:
         raise ValueError(f"Invalid display mode: {display_mode}")