RFC: streamlit_folium_vnext — Streamlit Components v2 rewrite

streamlit_folium_vnext/__init__.py

@@ -0,0 +1,3 @@
+from .api import FoliumResult, compile_folium, st_folium_vnext, st_leaflet
+
+__all__ = ["FoliumResult", "compile_folium", "st_folium_vnext", "st_leaflet"]

streamlit_folium_vnext/api.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any
+
+import folium
+import streamlit as st
+
+from streamlit_folium_vnext.compiler import compile_folium_map
+from streamlit_folium_vnext.component import mount_leaflet_component
+from streamlit_folium_vnext.models.spec import MapSpec
+
+_DEFAULT_ACCUMULATORS: dict[str, dict[str, Any]] = {
+    "click": {
+        "key": "clicks",
+        "initial": [],
+        "mode": "append",
+        "extract": lambda e: e["payload"],
+    },
+    "draw.created": {
+        "key": "drawn_features",
+        "initial": [],
+        "mode": "append",
+        "extract": lambda e: e["payload"].get("geojson"),
+    },
+    "draw.edited": {
+        "key": "drawn_features",
+        "initial": [],
+        "mode": "replace",
+        "extract": lambda e: e["payload"].get("features", []),
+    },
+    "draw.deleted": {
+        "key": "drawn_features",
+        "initial": [],
+        "mode": "replace",
+        "extract": lambda e: e["payload"].get("features", []),
+    },
+}
+
+
+def _session_key(component_key: str | None) -> str:
+    return f"_stf_vnext_state_{component_key or 'default'}"
+
+
+def _accumulate(
+    component_key: str | None, event: dict[str, Any] | None, state_spec: dict[str, Any]
+) -> dict[str, Any]:
+    sk = _session_key(component_key)
+    if sk not in st.session_state:
+        st.session_state[sk] = deepcopy(state_spec)
+
+    acc: dict[str, Any] = st.session_state[sk]
+
+    if event is None or not isinstance(event, dict):
+        return acc
+
+    event_type = event.get("type", "")
+    cfg = _DEFAULT_ACCUMULATORS.get(event_type)
+    if cfg is None:
+        return acc
+
+    target_key = cfg["key"]
+    if target_key not in acc:
+        return acc
+
+    value = cfg["extract"](event)
+    if cfg["mode"] == "append":
+        acc[target_key] = acc[target_key] + [value]
+    elif cfg["mode"] == "replace":
+        acc[target_key] = list(value) if isinstance(value, list) else [value]
+
+    st.session_state[sk] = acc
+    return acc
+
+
+class FoliumResult:
+    def __init__(self, raw_result: Any, accumulated: dict[str, Any] | None = None):
+        self._raw = raw_result
+        self._accumulated = accumulated or {}
+
+    def __getattr__(self, name: str) -> Any:
+        if name.startswith("_"):
+            raise AttributeError(name)
+        if name == "state":
+            return self._accumulated
+        return getattr(self._raw, name)
+
+    def __bool__(self) -> bool:
+        return bool(self._raw)
+
+    def __repr__(self) -> str:
+        return f"FoliumResult(raw={self._raw!r}, state={self._accumulated!r})"
+
+
+def compile_folium(obj: folium.Map, *, subscribe: list[str] | None = None) -> MapSpec:
+    return compile_folium_map(obj, subscribe=subscribe)
+
+
+def st_leaflet(
+    spec: MapSpec,
+    *,
+    key: str | None = None,
+    height: int = 500,
+    width: str = "stretch",
+    state: dict[str, Any] | None = None,
+):
+    d = spec.to_dict()
+    if key is not None:
+        d["map"]["id"] = key
+    raw = mount_leaflet_component(
+        spec=d,
+        key=key,
+        height=height,
+        width=width,
+    )
+    if state is not None:
+        event = getattr(raw, "event", None)
+        acc = _accumulate(key, event, state)
+        return FoliumResult(raw, acc)
+    return raw
+
+
+def st_folium_vnext(
+    obj: folium.Map,
+    *,
+    key: str | None = None,
+    height: int = 500,
+    width: str = "stretch",
+    subscribe: list[str] | None = None,
+    state: dict[str, Any] | None = None,
+):
+    spec = compile_folium(obj, subscribe=subscribe)
+    return st_leaflet(spec, key=key, height=height, width=width, state=state)

streamlit_folium_vnext/compiler/__init__.py

@@ -0,0 +1,3 @@
+from .compile_map import compile_folium_map
+
+__all__ = ["compile_folium_map"]

streamlit_folium_vnext/compiler/compile_map.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import folium
+import folium.plugins
+
+from streamlit_folium_vnext.compiler.context import CompileContext
+from streamlit_folium_vnext.compiler.plugins import (
+    compile_draw,
+    compile_feature_group,
+    compile_layer_control,
+)
+from streamlit_folium_vnext.compiler.plugins.map_objects import (
+    compile_circle,
+    compile_circle_marker,
+    compile_geojson,
+    compile_heat,
+    compile_marker,
+    compile_marker_cluster,
+    compile_polygon,
+    compile_polyline,
+    compile_tile_layer,
+)
+from streamlit_folium_vnext.compiler.registry import CompilerRegistry
+from streamlit_folium_vnext.models.spec import MapSpec
+
+registry = CompilerRegistry()
+registry.register(folium.raster_layers.TileLayer, compile_tile_layer)
+registry.register(folium.map.Marker, compile_marker)
+registry.register(folium.vector_layers.CircleMarker, compile_circle_marker)
+registry.register(folium.vector_layers.Circle, compile_circle)
+registry.register(folium.vector_layers.PolyLine, compile_polyline)
+registry.register(folium.vector_layers.Polygon, compile_polygon)
+registry.register(folium.features.GeoJson, compile_geojson)
+registry.register(folium.FeatureGroup, compile_feature_group)
+registry.register(folium.LayerControl, compile_layer_control)
+registry.register(folium.plugins.Draw, compile_draw)
+registry.register(folium.plugins.MarkerCluster, compile_marker_cluster)
+registry.register(folium.plugins.HeatMap, compile_heat)
+
+
+def compile_folium_map(
+    obj: folium.Map, *, subscribe: list[str] | None = None
+) -> MapSpec:
+    context = CompileContext()
+    spec = MapSpec(
+        map={
+            "id": context.allocate_id("map"),
+            "center": list(obj.location),
+            "zoom": obj.options.get("zoom") or obj.options.get("zoomStart"),
+            "options": dict(obj.options),
+        },
+        subscriptions=subscribe or [],
+    )
+
+    for child in obj._children.values():
+        handler = registry.resolve(child)
+        if handler is None:
+            continue
+        node = handler(child, context)
+        match node.kind:
+            case (
+                "tile_layer"
+                | "marker"
+                | "circle_marker"
+                | "circle"
+                | "polyline"
+                | "polygon"
+                | "geojson"
+                | "feature_group"
+                | "marker_cluster"
+                | "heat"
+            ):
+                spec.layers.append(node)
+            case "layer_control":
+                spec.controls.append(node)
+            case _:
+                spec.plugins.append(node)
+
+    return spec

streamlit_folium_vnext/compiler/context.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class CompileContext:
+    next_id: int = 0
+
+    def allocate_id(self, prefix: str) -> str:
+        self.next_id += 1
+        return f"{prefix}-{self.next_id}"

streamlit_folium_vnext/compiler/nodes.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from streamlit_folium_vnext.models.spec import MapNode
+
+
+def make_node(kind: str, node_id: str, **props) -> MapNode:
+    return MapNode(kind=kind, id=node_id, props=props)

streamlit_folium_vnext/compiler/plugins/__init__.py

@@ -0,0 +1,19 @@
+from .draw import compile_draw
+from .feature_group import compile_feature_group
+from .layer_control import compile_layer_control
+from .map_objects import (
+    compile_circle_marker,
+    compile_geojson,
+    compile_marker,
+    compile_tile_layer,
+)
+
+__all__ = [
+    "compile_circle_marker",
+    "compile_draw",
+    "compile_feature_group",
+    "compile_geojson",
+    "compile_layer_control",
+    "compile_marker",
+    "compile_tile_layer",
+]

streamlit_folium_vnext/compiler/plugins/draw.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import folium.plugins
+
+from streamlit_folium_vnext.compiler.context import CompileContext
+from streamlit_folium_vnext.compiler.nodes import make_node
+from streamlit_folium_vnext.models.spec import MapNode
+
+
+def compile_draw(obj: folium.plugins.Draw, context: CompileContext) -> MapNode:
+    return make_node(
+        "draw", context.allocate_id("draw"), options=getattr(obj, "options", {})
+    )

streamlit_folium_vnext/compiler/plugins/feature_group.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+import folium
+
+from streamlit_folium_vnext.compiler.context import CompileContext
+from streamlit_folium_vnext.compiler.nodes import make_node
+from streamlit_folium_vnext.models.spec import MapNode
+
+
+def compile_feature_group(obj: folium.FeatureGroup, context: CompileContext) -> MapNode:
+    return make_node(
+        "feature_group",
+        context.allocate_id("feature-group"),
+        name=getattr(obj, "layer_name", None),
+    )

streamlit_folium_vnext/compiler/plugins/layer_control.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import folium
+
+from streamlit_folium_vnext.compiler.context import CompileContext
+from streamlit_folium_vnext.compiler.nodes import make_node
+from streamlit_folium_vnext.models.spec import MapNode
+
+
+def compile_layer_control(obj: folium.LayerControl, context: CompileContext) -> MapNode:
+    options = dict(obj.options) if hasattr(obj, "options") else {}
+    return make_node(
+        "layer_control",
+        context.allocate_id("layer-control"),
+        position=options.get("position", "topright"),
+        options=options,
+    )