Ability to provide feature group as list (#165)

* fix: actually remove the feature group when needed * feat: possibility to provide the feature group as a list * fix: changed expected type and documentation * fix title * added test to check that the feature groups are appearing and disappearing as requested * Linting and expand test a bit * update RELEASE * Tsx linting --------- Co-authored-by: Zachary Blackwood <[email protected]>
randyzwitch · Dec 12, 2023 · 32566f6 · 32566f6
1 parent 884746d
commit 32566f6
Show file tree

Hide file tree

Showing 5 changed files with 183 additions and 35 deletions.
diff --git a/examples/pages/dynamic_updates.py b/examples/pages/dynamic_updates.py
@@ -4,8 +4,10 @@
 
 import folium
 import folium.features
+import geopandas as gpd
 import pandas as pd
 import requests
+import shapely
 import streamlit as st
 
 from streamlit_folium import st_folium
@@ -19,8 +21,8 @@
 "# Dynamic Updates -- Click on a marker"
 
 st.subheader(
-    """Use new arguments `center`, `zoom`, and `feature_group_to_add` to update the map
-    without re-rendering it."""
+    "Use new arguments `center`, `zoom`, and `feature_group_to_add` to update the map "
+    "without re-rendering it."
 )
 
 
@@ -105,6 +107,76 @@ def main():
         st.session_state["selected_state"] = state
         st.experimental_rerun()
 
+    st.write("## Dynamic feature group updates")
+
+    START_LOCATION = [37.7944347109497, -122.398077892527]
+    START_ZOOM = 17
+
+    if "feature_group" not in st.session_state:
+        st.session_state["feature_group"] = None
+
+    wkt1 = (
+        "POLYGON ((-122.399077892527 37.7934347109497, -122.398922660838 "
+        "37.7934544916178, -122.398980265018 37.7937266504805, -122.399133972495 "
+        "37.7937070646238, -122.399077892527 37.7934347109497))"
+    )
+    wkt2 = (
+        "POLYGON ((-122.397416 37.795017, -122.397137 37.794712, -122.396332 37.794983,"
+        " -122.396171 37.795483, -122.396858 37.795695, -122.397652 37.795466, "
+        "-122.397759 37.79511, -122.397416 37.795017))"
+    )
+
+    polygon_1 = shapely.wkt.loads(wkt1)
+    polygon_2 = shapely.wkt.loads(wkt2)
+
+    gdf1 = gpd.GeoDataFrame(geometry=[polygon_1]).set_crs(epsg=4326)
+    gdf2 = gpd.GeoDataFrame(geometry=[polygon_2]).set_crs(epsg=4326)
+
+    style_parcels = {
+        "fillColor": "#1100f8",
+        "color": "#1100f8",
+        "fillOpacity": 0.13,
+        "weight": 2,
+    }
+    style_buildings = {
+        "color": "#ff3939",
+        "fillOpacity": 0,
+        "weight": 3,
+        "opacity": 1,
+        "dashArray": "5, 5",
+    }
+
+    polygon_folium1 = folium.GeoJson(data=gdf1, style_function=lambda x: style_parcels)
+    polygon_folium2 = folium.GeoJson(
+        data=gdf2, style_function=lambda x: style_buildings
+    )
+
+    map = folium.Map(
+        location=START_LOCATION,
+        zoom_start=START_ZOOM,
+        tiles="OpenStreetMap",
+        max_zoom=21,
+    )
+
+    fg1 = folium.FeatureGroup(name="Parcels")
+    fg1.add_child(polygon_folium1)
+
+    fg2 = folium.FeatureGroup(name="Buildings")
+    fg2.add_child(polygon_folium2)
+
+    fg_dict = {"Parcels": fg1, "Buildings": fg2, "None": None, "Both": [fg1, fg2]}
+
+    fg = st.radio("Feature Group", ["Parcels", "Buildings", "None", "Both"])
+
+    st_folium(
+        map,
+        width=800,
+        height=450,
+        returned_objects=[],
+        feature_group_to_add=fg_dict[fg],
+        debug=True,
+    )
+
 
 if __name__ == "__main__":
     main()
diff --git a/streamlit_folium/__init__.py b/streamlit_folium/__init__.py
@@ -5,7 +5,7 @@
 import re
 import warnings
 from textwrap import dedent
-from typing import Dict, Iterable, List
+from typing import Iterable
 
 import branca
 import folium
@@ -154,23 +154,26 @@ def _get_map_string(fig: folium.Map) -> str:
 def _get_feature_group_string(
     feature_group_to_add: folium.FeatureGroup,
     map: folium.Map,
+    idx: int = 0,
 ) -> str:
-    feature_group_to_add._id = "feature_group"
+    feature_group_to_add._id = f"feature_group_{idx}"
     feature_group_to_add.add_to(map)
     feature_group_to_add.render()
     feature_group_string = generate_leaflet_string(
-        feature_group_to_add, base_id="feature_group"
+        feature_group_to_add, base_id=f"feature_group_{idx}"
     )
     m_id = get_full_id(map)
     feature_group_string = feature_group_string.replace(m_id, "map_div")
     feature_group_string = dedent(feature_group_string)
 
     feature_group_string += dedent(
-        """
-        map_div.addLayer(feature_group_feature_group);
-        window.feature_group = feature_group_feature_group;
+        f"""
+        map_div.addLayer(feature_group_feature_group_{idx});
+        window.feature_group = window.feature_group || [];
+        window.feature_group.push(feature_group_feature_group_{idx});
         """
     )
+
     return feature_group_string
 
 
@@ -182,7 +185,7 @@ def st_folium(
     returned_objects: Iterable[str] | None = None,
     zoom: int | None = None,
     center: tuple[float, float] | None = None,
-    feature_group_to_add: folium.FeatureGroup | None = None,
+    feature_group_to_add: list[folium.FeatureGroup] | folium.FeatureGroup | None = None,
     return_on_hover: bool = False,
     use_container_width: bool = False,
     debug: bool = False,
@@ -212,7 +215,7 @@ def st_folium(
         The center of the map. If None, the center will be set to the default
         center of the map. NOTE that if this center is changed, it will *not* reload
         the map, but simply dynamically change the center.
-    feature_group_to_add: folium.FeatureGroup or None
+    feature_group_to_add: List[folium.FeatureGroup] or folium.FeatureGroup or None
         If you want to dynamically add features to a feature group, you can pass
         the feature group here. NOTE that if you add a feature to the map, it
         will *not* reload the map, but simply dynamically add the feature.
@@ -257,7 +260,7 @@ def st_folium(
 
     m_id = get_full_id(folium_map)
 
-    def bounds_to_dict(bounds_list: List[List[float]]) -> Dict[str, Dict[str, float]]:
+    def bounds_to_dict(bounds_list: list[list[float]]) -> dict[str, dict[str, float]]:
         southwest, northeast = bounds_list
         return {
             "_southWest": {
@@ -302,10 +305,15 @@ def bounds_to_dict(bounds_list: List[List[float]]) -> Dict[str, Dict[str, float]
     # on the frontend.
     feature_group_string = None
     if feature_group_to_add is not None:
-        feature_group_string = _get_feature_group_string(
-            feature_group_to_add,
-            map=folium_map,
-        )
+        if isinstance(feature_group_to_add, folium.FeatureGroup):
+            feature_group_to_add = [feature_group_to_add]
+        feature_group_string = ""
+        for idx, feature_group in enumerate(feature_group_to_add):
+            feature_group_string += _get_feature_group_string(
+                feature_group,
+                map=folium_map,
+                idx=idx,
+            )
 
     if debug:
         with st.expander("Show generated code"):

diff --git a/streamlit_folium/frontend/src/index.tsx b/streamlit_folium/frontend/src/index.tsx
@@ -1,3 +1,4 @@
+import { Layer } from "leaflet"
 import { RenderData, Streamlit } from "streamlit-component-lib"
 import { debounce } from "underscore"
 import { circleToPolygon } from "./circle-to-polygon"
@@ -239,28 +240,31 @@ function onRender(event: Event): void {
     }
   }
 
-  if (
-    feature_group &&
-    feature_group !== window.__GLOBAL_DATA__.last_feature_group
-  ) {
-    if (window.feature_group) {
-      window.map.removeLayer(window.feature_group)
+  if (feature_group !== window.__GLOBAL_DATA__.last_feature_group) {
+    if (window.feature_group && window.feature_group.length > 0) {
+      window.feature_group.forEach((layer: Layer) => {
+        window.map.removeLayer(layer)
+      })
     }
-    // Though using `eval` is generally a bad idea, we're using it here
-    // because we're evaluating code that we've generated ourselves on the
-    // Python side. This is safe because we're not evaluating user input, so this
-    // couldn't be used to execute arbitrary code.
 
-    // eslint-disable-next-line
-    eval(feature_group)
     window.__GLOBAL_DATA__.last_feature_group = feature_group
-    for (let key in window.map._layers) {
-      let layer = window.map._layers[key]
-      layer.off("click", onLayerClick)
-      layer.on("click", onLayerClick)
-      if (return_on_hover) {
-        layer.off("mouseover", onLayerClick)
-        layer.on("mouseover", onLayerClick)
+
+    if (feature_group) {
+      // Though using `eval` is generally a bad idea, we're using it here
+      // because we're evaluating code that we've generated ourselves on the
+      // Python side. This is safe because we're not evaluating user input, so this
+      // couldn't be used to execute arbitrary code.
+
+      // eslint-disable-next-line
+      eval(feature_group)
+      for (let key in window.map._layers) {
+        let layer = window.map._layers[key]
+        layer.off("click", onLayerClick)
+        layer.on("click", onLayerClick)
+        if (return_on_hover) {
+          layer.off("mouseover", onLayerClick)
+          layer.on("mouseover", onLayerClick)
+        }
       }
     }
   }

diff --git a/tests/test_frontend.py b/tests/test_frontend.py
@@ -268,3 +268,67 @@ def test_grouped_layer_control(page: Page):
     page.frame_locator('iframe[title="streamlit_folium\\.st_folium"]').get_by_label(
         "g2"
     ).check()
+
+
+def test_dynamic_feature_group_update(page: Page):
+    page.get_by_role("link", name="dynamic updates").click()
+    page.get_by_text("Show generated code").click()
+
+    # Test showing only Parcel layer
+    page.get_by_test_id("stRadio").get_by_text("Parcels").click()
+    expect(
+        page.frame_locator('iframe[title="streamlit_folium\\.st_folium"] >> nth=1')
+        .locator("path")
+        .first
+    ).to_be_visible()
+    expect(
+        page.frame_locator('iframe[title="streamlit_folium\\.st_folium"] >> nth=1')
+        .locator("path")
+        .nth(1)
+    ).to_be_hidden()
+    expect(
+        page.get_by_text('"fillColor"')
+    ).to_be_visible()  # fillColor only present in parcel style
+    expect(
+        page.get_by_text('"dashArray"')
+    ).to_be_hidden()  # dashArray only present in building style
+
+    # Test showing only Building layer
+    page.get_by_test_id("stRadio").get_by_text("Buildings").click()
+    expect(
+        page.frame_locator('iframe[title="streamlit_folium\\.st_folium"] >> nth=1')
+        .locator("path")
+        .first
+    ).to_be_visible()
+    expect(
+        page.frame_locator('iframe[title="streamlit_folium\\.st_folium"] >> nth=1')
+        .locator("path")
+        .nth(1)
+    ).to_be_hidden()
+    expect(page.get_by_text("fillColor")).to_be_hidden()
+    expect(page.get_by_text("dashArray")).to_be_visible()
+
+    # Test showing no layers
+    page.get_by_test_id("stRadio").get_by_text("None").click()
+    expect(
+        page.frame_locator('iframe[title="streamlit_folium\\.st_folium"] >> nth=1')
+        .locator("path")
+        .first
+    ).to_be_hidden()
+    expect(page.get_by_text("fillColor")).to_be_hidden()
+    expect(page.get_by_text("dashArray")).to_be_hidden()
+
+    # Test showing both layers
+    page.get_by_test_id("stRadio").get_by_text("Both").click()
+    expect(
+        page.frame_locator('iframe[title="streamlit_folium\\.st_folium"] >> nth=1')
+        .locator("path")
+        .first
+    ).to_be_visible()
+    expect(
+        page.frame_locator('iframe[title="streamlit_folium\\.st_folium"] >> nth=1')
+        .locator("path")
+        .nth(1)
+    ).to_be_visible()
+    expect(page.get_by_text("fillColor")).to_be_visible()
+    expect(page.get_by_text("dashArray")).to_be_visible()
diff --git a/tests/test_package.py b/tests/test_package.py
@@ -89,7 +89,7 @@ def test_feature_group():
 
     fg_str = _get_feature_group_string(fg, m)
 
-    assert "var feature_group_feature_group = L.featureGroup(" in fg_str
+    assert "var feature_group_feature_group_0 = L.featureGroup(" in fg_str
     assert ".addTo(map_div);" in fg_str