Add leaflet_method decorator (#2157)

* Add leaflet_method decorator

This allows javascript method calls to be generated using
an empty method body and a signature. Useful for writing plugins.

As an example (and the driving motivation) I implemented this on
the geoman plugin.

* Fix tests

* Implement all useful leaflet methods on geoman

Only the free modules for now

* Add a snapshot test

Author: hansthen

folium/elements.py

@@ -1,3 +1,4 @@
+from functools import wraps
 from typing import List, Tuple
 
 from branca.element import (
@@ -9,7 +10,15 @@
 )
 
 from folium.template import Template
-from folium.utilities import JsCode
+from folium.utilities import JsCode, camelize
+
+
+def leaflet_method(fn):
+    @wraps(fn)
+    def inner(self, *args, **kwargs):
+        self.add_child(MethodCall(self, fn.__name__, *args, **kwargs))
+
+    return inner
 
 
 class JSCSSMixin(MacroElement):
@@ -148,3 +157,27 @@ def __init__(self, element_name: str, element_parent_name: str):
         super().__init__()
         self.element_name = element_name
         self.element_parent_name = element_parent_name
+
+
+class MethodCall(MacroElement):
+    """Abstract class to add an element to another element."""
+
+    _template = Template(
+        """
+        {% macro script(this, kwargs) %}
+            {{ this.target }}.{{ this.method }}(
+                {% for arg in this.args %}
+                    {{ arg | tojavascript }},
+                {% endfor %}
+                {{ this.kwargs | tojavascript }}
+            );
+        {% endmacro %}
+    """
+    )
+
+    def __init__(self, target: MacroElement, method: str, *args, **kwargs):
+        super().__init__()
+        self.target = target.get_name()
+        self.method = camelize(method)
+        self.args = args
+        self.kwargs = kwargs

folium/plugins/geoman.py

@@ -1,6 +1,6 @@
 from branca.element import MacroElement
 
-from folium.elements import JSCSSMixin
+from folium.elements import JSCSSMixin, leaflet_method
 from folium.template import Template
 from folium.utilities import remove_empty
 
@@ -22,6 +22,8 @@ class GeoMan(JSCSSMixin, MacroElement):
     _template = Template(
         """
         {% macro script(this, kwargs) %}
+            /* ensure the name is usable */
+            var {{this.get_name()}} = {{this._parent.get_name()}}.pm;
             {%- if this.feature_group  %}
                 var drawnItems_{{ this.get_name() }} =
                     {{ this.feature_group.get_name() }};
@@ -32,12 +34,12 @@ class GeoMan(JSCSSMixin, MacroElement):
                         {{ this._parent.get_name() }}
                     );
             {%- endif %}
-            /* The global varianble below is needed to prevent streamlit-folium
+            /* The global variable below is needed to prevent streamlit-folium
                from barfing :-(
             */
             var drawnItems = drawnItems_{{ this.get_name() }};
 
-            {{this._parent.get_name()}}.pm.addControls(
+            {{this.get_name()}}.addControls(
                 {{this.options|tojavascript}}
             )
             drawnItems_{{ this.get_name() }}.eachLayer(function(layer){
@@ -60,12 +62,6 @@ class GeoMan(JSCSSMixin, MacroElement):
                     {{handler}}
                 );
                 {%- endfor %}
-                drawnItems_{{ this.get_name() }}.addLayer(layer);
-            });
-            {{ this._parent.get_name() }}.on("pm:remove", function(e) {
-                var layer = e.layer,
-                    type = e.layerType;
-                drawnItems_{{ this.get_name() }}.removeLayer(layer);
             });
 
         {% endmacro %}
@@ -85,17 +81,65 @@ class GeoMan(JSCSSMixin, MacroElement):
         )
     ]
 
-    def __init__(
-        self,
-        position="topleft",
-        feature_group=None,
-        on=None,
-        **kwargs,
-    ):
+    def __init__(self, position="topleft", feature_group=None, on=None, **kwargs):
         super().__init__()
         self._name = "GeoMan"
         self.feature_group = feature_group
         self.on = on or {}
-        self.options = remove_empty(
-            position=position, layer_group=feature_group, **kwargs
-        )
+        self.options = remove_empty(position=position, **kwargs)
+
+    @leaflet_method
+    def set_global_options(self, **kwargs):
+        pass
+
+    @leaflet_method
+    def enable_draw(self, shape, /, **kwargs):
+        pass
+
+    @leaflet_method
+    def disable_draw(self):
+        pass
+
+    @leaflet_method
+    def set_path_options(self, *, options_modifier, **options):
+        pass
+
+    @leaflet_method
+    def enable_global_edit_mode(self, **options):
+        pass
+
+    @leaflet_method
+    def disable_global_edit_mode(self):
+        pass
+
+    @leaflet_method
+    def enable_global_drag_mode(self):
+        pass
+
+    @leaflet_method
+    def disable_global_drag_mode(self):
+        pass
+
+    @leaflet_method
+    def enable_global_removal_mode(self):
+        pass
+
+    @leaflet_method
+    def disable_global_removal_mode(self):
+        pass
+
+    @leaflet_method
+    def enable_global_cut_mode(self):
+        pass
+
+    @leaflet_method
+    def disable_global_cut_mode(self):
+        pass
+
+    @leaflet_method
+    def enable_global_rotation_mode(self):
+        pass
+
+    @leaflet_method
+    def disable_global_rotation_mode(self):
+        pass

tests/plugins/test_geoman.py

@@ -20,7 +20,7 @@ def test_geoman():
     # the map
     tmpl = Template(
         """
-        {{this._parent.get_name()}}.pm.addControls(
+        {{this.get_name()}}.addControls(
             {{this.options|tojavascript}}
         )
     """

tests/snapshots/modules/geoman_customizations.py

@@ -0,0 +1,62 @@
+import folium
+from folium import JsCode
+from folium.plugins import GeoMan, MousePosition
+
+m = folium.Map(tiles=None, location=[39.949610, -75.150282], zoom_start=5)
+MousePosition().add_to(m)
+
+# This can be used to test the connection to streamlit
+# by returning the resulting GeoJson
+handler = JsCode(
+    """
+  (e) => {
+      var map = %(map)s;
+      var layers = L.PM.Utils.findLayers(map);
+      var lg = L.layerGroup(layers);
+      console.log(lg.toGeoJSON());
+  }
+  """  # noqa: UP031
+    % dict(map=m.get_name())
+)
+
+# For manual testing
+click = JsCode(
+    """
+  (e) => {
+      console.log(e.target);
+      console.log(e.target.toGeoJSON());
+  }
+  """
+)
+
+# Just a few customizations for the snapshot tests
+# The test succeeds if the position is to the right
+# and if the buttons for markers and circles are not
+# shown.
+gm = GeoMan(
+    position="topright", draw_marker=False, draw_circle=False, on={"click": click}
+).add_to(m)
+
+# For manual testing of the global options
+gm.set_global_options(
+    {
+        "snappable": True,
+        "snapDistance": 20,
+    }
+)
+
+# Make rectangles green
+gm.enable_draw("Rectangle", path_options={"color": "green"})
+gm.disable_draw()
+
+# On any event that updates the layers, we trigger the handler
+event_handlers = {
+    "pm:create": handler,
+    "pm:remove": handler,
+    "pm:update": handler,
+    "pm:rotateend": handler,
+    "pm:cut": handler,
+    "pm:undoremove": handler,
+}
+
+m.on(**event_handlers)