feat: support marimo poc version (#647)

* feat: support marimo poc version

* chore: bump to v0.4.9.11

Author: longxiaofei

app/package.json

@@ -15,6 +15,7 @@
     "serve": "vite preview"
   },
   "dependencies": {
+    "@anywidget/react": "^0.0.8",
     "@headlessui/react": "^1.7.14",
     "@heroicons/react": "^2.0.8",
     "@kanaries/graphic-walker": "0.4.70",

app/src/index.tsx

@@ -7,14 +7,15 @@ import type { VizSpecStore } from '@kanaries/graphic-walker/store/visualSpecStor
 import type { IGWHandler, IViewField, ISegmentKey, IDarkMode, IChatMessage, IRow } from '@kanaries/graphic-walker/interfaces';
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { Streamlit, withStreamlitConnection } from "streamlit-component-lib"
+import { createRender, useModel } from "@anywidget/react";
 
 import Options from './components/options';
 import { IAppProps } from './interfaces';
 
 import { loadDataSource, postDataService, finishDataService, getDatasFromKernelBySql, getDatasFromKernelByPayload } from './dataSource';
 
 import commonStore from "./store/common";
-import { initJupyterCommunication, initHttpCommunication, streamlitComponentCallback } from "./utils/communication";
+import { initJupyterCommunication, initHttpCommunication, streamlitComponentCallback, initAnywidgetCommunication } from "./utils/communication";
 import communicationStore from "./store/communication"
 import { setConfig } from './utils/userConfig';
 import CodeExportModal from './components/codeExportModal';
@@ -209,7 +210,7 @@ const ExploreApp: React.FC<IAppProps & {initChartFlag: boolean}> = (props) => {
     const exportTool = getExportTool(setExportOpen);
 
     const tools = [exportTool];
-    if ((props.env === "jupyter_widgets" || props.env === "streamlit" || props.env === "gradio") && props.useSaveTool) {
+    if (props.env && ["jupyter_widgets", "streamlit", "gradio", "marimo"].indexOf(props.env) !== -1 && props.useSaveTool) {
         const saveTool = getSaveTool(props, gwRef, storeRef, isChanged, setIsChanged);
         tools.push(saveTool);
     }
@@ -368,6 +369,16 @@ const initOnHttpCommunication = async(props: IAppProps) => {
     await initDslParser();
 }
 
+const initOnAnywidgetCommunication = async(props: IAppProps, model: import("@anywidget/types").AnyModel) => {
+    const comm = await initAnywidgetCommunication(props.id, model);
+    communicationStore.setComm(comm);
+    if ((props.gwMode === "explore" || props.gwMode === "filter_renderer") && props.needLoadLastSpec) {
+        const visSpecResp = await comm.sendMsg("get_latest_vis_spec", {});
+        props.visSpec = visSpecResp["data"]["visSpec"];
+    }
+    await initDslParser();
+}
+
 const defaultInit = async(props: IAppProps) => {}
 
 function GWalkerComponent(props: IAppProps) {
@@ -568,4 +579,29 @@ const StreamlitGWalker = () => {
     )
 }
 
-export default { GWalker, PreviewApp, ChartPreviewApp, StreamlitGWalker }
+function AnywidgetGWalkerApp() {
+    const [inited, setInited] = useState(false);
+    const model = useModel();
+    const props = JSON.parse(model.get("props")) as IAppProps;
+    props.visSpec = FormatSpec(props.visSpec, props.rawFields);
+
+    useEffect(() => {
+        initOnAnywidgetCommunication(props, model).then(() => {
+            setInited(true);
+        })
+    }, []);
+
+    return (
+        <React.StrictMode>
+            {!inited && <div>Loading...</div>}
+            {inited && (
+                <MainApp darkMode={props.dark}>
+                    <GWalkerComponent {...props} />
+                </MainApp>
+            )}
+        </React.StrictMode>
+    );
+}
+
+
+export default { GWalker, PreviewApp, ChartPreviewApp, StreamlitGWalker, render: createRender(AnywidgetGWalkerApp) }

app/src/utils/communication.tsx

@@ -255,5 +255,63 @@ const streamlitComponentCallback = (data: any) => {
     }
 }
 
+const initAnywidgetCommunication = async(gid: string, model: import("@anywidget/types").AnyModel) => {
+    const bufferMap = new Map<string, any>();
+
+    const onMessage = (msg: string) => {
+        const data = JSON.parse(msg);
+        const action = data.action;
+        if (action === "finish_request") {
+            bufferMap.set(data.rid, data.data);
+            document.dispatchEvent(new CustomEvent(getSignalName(data.rid)));
+            return
+        }
+    }
+
+    model.on("msg:custom", msg => {
+        if (msg.type !== "pyg_response") {
+            return;
+        }
+        onMessage(msg.data);
+    });
+
+    const sendMsg = async(action: string, data: any, timeout: number = 30_000) => {
+        const rid = uuidv4();
+        const promise = new Promise<any>((resolve, reject) => {
+            setTimeout(() => {
+                sendMsgAsync(action, data, rid);
+            }, 0);
+            const timer = setTimeout(() => {
+                raiseRequestError("communication timeout", 0);
+                reject(new Error("get result timeout"));
+            }, timeout);
+            document.addEventListener(getSignalName(rid), (_) => {
+                clearTimeout(timer);
+                const resp = bufferMap.get(rid);
+                if (resp.code !== 0) {
+                    raiseRequestError(resp.message, resp.code);
+                    reject(new Error(resp.message));
+                }
+                resolve(resp);
+            });
+        });
+
+        return promise;
+    }
+
+    const sendMsgAsync = (action: string, data: any, rid: string | null) => {
+        rid = rid ?? uuidv4();
+        model.send({type: "pyg_request", msg: { gid, rid, action, data }});
+    }
+
+    const registerEndpoint = (_: string, __: (data: any) => any) => {}
+
+    return {
+        sendMsg,
+        registerEndpoint,
+        sendMsgAsync,
+    }
+}
+
 export type { ICommunication };
-export { initJupyterCommunication, initHttpCommunication, streamlitComponentCallback };
+export { initJupyterCommunication, initHttpCommunication, streamlitComponentCallback, initAnywidgetCommunication };

app/vite.config.ts

@@ -17,7 +17,7 @@ export default defineConfig((config: ConfigEnv) => {
         entry: path.resolve(__dirname, './src/index.tsx'),
         name: 'PyGWalkerApp',
         fileName: (format) => `pygwalker-app.${format}.js`,
-        formats: ['iife']
+        formats: ['iife', "es"]
       },
       minify: 'esbuild',
       sourcemap: false,
@@ -77,6 +77,8 @@ export default defineConfig((config: ConfigEnv) => {
       rollupOptions: {
         external: modulesNotToBundle,
         output: {
+          manualChunks: undefined,
+          inlineDynamicImports: true,
           globals: {
             'react': 'React',
             'react-dom': 'ReactDOM',

app/yarn.lock

@@ -15,6 +15,18 @@
     "@jridgewell/gen-mapping" "^0.3.0"
     "@jridgewell/trace-mapping" "^0.3.9"
 
+"@anywidget/react@^0.0.8":
+  version "0.0.8"
+  resolved "https://registry.yarnpkg.com/@anywidget/react/-/react-0.0.8.tgz#64b73cf9c9bad7bad180c8ba44134f09cc6eb44f"
+  integrity sha512-obr4EasXgWra485u+G4V3Msn7A1EOnowarvR62FRjpv2Rz6AyOoLMz2B03Z9j3DrWdD0634fMGu5ZAeRyjuV4w==
+  dependencies:
+    "@anywidget/types" "^0.2.0"
+
+"@anywidget/types@^0.2.0":
+  version "0.2.0"
+  resolved "https://registry.yarnpkg.com/@anywidget/types/-/types-0.2.0.tgz#6bae4e4fa36d193f565b0b78dc7eac50bb71beb4"
+  integrity sha512-+XtK4uwxRd4JpuevUMhirrbvC0V4yCA/i0lEjhmSAtOaxiXIg/vBKzaSonDuoZ1a9LEjUXTW2+m7w+ULgsJYvg==
+
 "@babel/code-frame@^7.0.0", "@babel/code-frame@^7.12.13", "@babel/code-frame@^7.22.13":
   version "7.22.13"
   resolved "https://registry.yarnpkg.com/@babel/code-frame/-/code-frame-7.22.13.tgz#e3c1c099402598483b7a8c46a721d1038803755e"

pygwalker/__init__.py

@@ -10,7 +10,7 @@
 from pygwalker.services.global_var import GlobalVarManager
 from pygwalker.services.kaggle import show_tips_user_kaggle as __show_tips_user_kaggle
 
-__version__ = "0.4.9.10"
+__version__ = "0.4.9.11"
 __hash__ = __rand_str()
 
 from pygwalker.api.jupyter import walk, render, table

pygwalker/api/marimo.py

@@ -0,0 +1,85 @@
+from typing import Union, List, Optional
+import inspect
+import json
+import pathlib
+
+from typing_extensions import Literal
+
+from .pygwalker import PygWalker
+from pygwalker.data_parsers.base import FieldSpec
+from pygwalker.data_parsers.database_parser import Connector
+from pygwalker._typing import DataFrame, IAppearance, IThemeKey
+from pygwalker.services.format_invoke_walk_code import get_formated_spec_params_code_from_frame
+from pygwalker.communications.anywidget_comm import AnywidgetCommunication
+import marimo as mo
+import anywidget
+import traitlets
+
+
+class _WalkerWidget(anywidget.AnyWidget):
+    """WalkerWidget"""
+    _esm = (pathlib.Path(__file__).parent.parent / "templates" / "dist" / "pygwalker-app.es.js").read_text()
+    props = traitlets.Unicode("").tag(sync=True)
+
+
+def walk(
+    dataset: Union[DataFrame, Connector, str],
+    gid: Union[int, str] = None,
+    *,
+    field_specs: Optional[List[FieldSpec]] = None,
+    theme_key: IThemeKey = 'g2',
+    appearance: IAppearance = 'media',
+    spec: str = "",
+    show_cloud_tool: bool = True,
+    kanaries_api_key: str = "",
+    default_tab: Literal["data", "vis"] = "vis",
+    **kwargs
+):
+    """Walk through pandas.DataFrame df with Graphic Walker
+
+    Args:
+        - dataset (pl.DataFrame | pd.DataFrame | Connector, optional): dataframe.
+        - gid (Union[int, str], optional): GraphicWalker container div's id ('gwalker-{gid}')
+
+    Kargs:
+        - field_specs (List[FieldSpec], optional): Specifications of some fields. They'll been automatically inferred from `df` if some fields are not specified.
+        - theme_key ('vega' | 'g2' | 'streamlit'): theme type.
+        - appearance (Literal['media' | 'light' | 'dark']): 'media': auto detect OS theme.
+        - spec (str): chart config data. config id, json, remote file url
+        - kanaries_api_key (str): kanaries api key, Default to "".
+        - default_tab (Literal["data", "vis"]): default tab to show. Default to "vis"
+    """
+    if field_specs is None:
+        field_specs = []
+
+    source_invoke_code = get_formated_spec_params_code_from_frame(
+        inspect.stack()[1].frame
+    )
+
+    widget = _WalkerWidget()
+    walker = PygWalker(
+        gid=gid,
+        dataset=dataset,
+        field_specs=field_specs,
+        spec=spec,
+        source_invoke_code=source_invoke_code,
+        theme_key=theme_key,
+        appearance=appearance,
+        show_cloud_tool=show_cloud_tool,
+        use_preview=False,
+        kernel_computation=True,
+        use_save_tool=True,
+        gw_mode="explore",
+        is_export_dataframe=True,
+        kanaries_api_key=kanaries_api_key,
+        default_tab=default_tab,
+        cloud_computation=False,
+        **kwargs
+    )
+    comm = AnywidgetCommunication(walker.gid)
+
+    widget.props = json.dumps(walker._get_props("marimo", []))
+    comm.register_widget(widget)
+    walker._init_callback(comm)
+
+    return mo.ui.anywidget(widget)

pygwalker/communications/anywidget_comm.py

@@ -0,0 +1,42 @@
+from typing import Any, Dict, Optional, List
+import uuid
+import json
+
+import anywidget
+
+from .base import BaseCommunication
+from pygwalker.utils.encode import DataFrameEncoder
+
+
+class AnywidgetCommunication(BaseCommunication):
+    """communication class for anywidget"""
+    def register_widget(self, widget: anywidget.AnyWidget) -> None:
+        """register widget"""
+        self.widget = widget
+        self.widget.on_msg(self._on_mesage)
+
+    def send_msg_async(self, action: str, data: Dict[str, Any], rid: Optional[str] = None):
+        """send message base on anywidget"""
+        if rid is None:
+            rid = uuid.uuid1().hex
+        msg = {
+            "gid": self.gid,
+            "rid": rid,
+            "action": action,
+            "data": data
+        }
+        self.widget.send({"type": "pyg_response", "data": json.dumps(msg, cls=DataFrameEncoder)})
+
+    def _on_mesage(self, _: anywidget.AnyWidget, data: Dict[str, Any], buffers: List[Any]):
+        if data.get("type", "") != "pyg_request":
+            return
+        
+        msg = data["msg"]
+        action = msg["action"]
+        rid = msg["rid"]
+
+        if action == "finish_request":
+            return
+
+        resp = self._receive_msg(action, msg["data"])
+        self.send_msg_async("finish_request", resp, rid)

pyproject.toml

@@ -36,6 +36,8 @@ dependencies = [
     "numpy<2.0.0",
     "ipylab<=1.0.0",
     "quickjs",
+    "traitlets",
+    "anywidget",
 ]
 [project.urls]
 homepage = "https://kanaries.net/pygwalker"