Fix test_marker_cluster

Author: Martin Journois

folium/plugins/marker_cluster.py

@@ -5,48 +5,66 @@
 
 Creates a MarkerCluster plugin to add on a folium map.
 """
-import json
+#import json
+from jinja2 import Template
 
-from .plugin import Plugin
+from folium.element import JavascriptLink, CssLink, MacroElement, Figure
+from folium.map import Popup, Icon, Marker
 
-class MarkerCluster(Plugin):
-    """Adds a MarkerCluster layer on the map."""
-    def __init__(self, data):
+class MarkerCluster(MacroElement):
+    def __init__(self, locations, popups=None, icons=None):
         """Creates a MarkerCluster plugin to append into a map with
-        Map.add_plugin.
+        Map.add_children.
 
         Parameters
         ----------
-            data: list of list or array of shape (n,3).
-                Data points of the form [[lat, lng, popup]].
+            locations: list of list or array of shape (n,2).
+                Data points of the form [[lat, lng]].
+
+            popups: list of length n.
+                Popup for each marker.
+
+            icons: list of length n.
+                Icon for each marker.
         """
         super(MarkerCluster, self).__init__()
-        self.plugin_name = 'MarkerCluster'
-        self.data = [tuple(x) for x in data]
-
-    def render_header(self, nb):
-        """Generates the HTML part of the plugin."""
-        return """
-        <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/leaflet.markercluster/0.4.0/MarkerCluster.css" />
-        <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/leaflet.markercluster/0.4.0/MarkerCluster.Default.css" />
-        <script src="https://cdnjs.cloudflare.com/ajax/libs/leaflet.markercluster/0.4.0/leaflet.markercluster.js"></script>
-        """ if nb==0 else ""
-
-    def render_js(self, nb):
-        """Generates the Javascript part of the plugin."""
-        out = """
-        var addressPoints = """+json.dumps(self.data)+""";
-
-        var markers = L.markerClusterGroup();
-
-        for (var i = 0; i < addressPoints.length; i++) {
-            var a = addressPoints[i];
-            var title = a[2];
-            var marker = L.marker(new L.LatLng(a[0], a[1]), { title: title });
-            marker.bindPopup(title);
-            markers.addLayer(marker);
-        }
-
-        map.addLayer(markers);
-        """
-        return out
+        self._name = 'MarkerCluster'
+
+        if popups is None:
+            popups = [None]*len(locations)
+        if icons is None:
+            icons = [None]*len(locations)
+
+        for location, popup, icon in zip(locations,popups,icons):
+            if popup is None or isinstance(popup, Popup):
+                p = popup
+            else:
+                p = Popup(popup)
+            if icon is None or isinstance(icon, Icon):
+                i = icon
+            else:
+                i = Icon(icon)
+            self.add_children(Marker(location, popup=p, icon=i))
+
+        self._template = Template(u"""
+            {% macro script(this, kwargs) %}
+                var {{this.get_name()}} = L.markerClusterGroup();
+                {{this._parent.get_name()}}.addLayer({{this.get_name()}});
+            {% endmacro %}
+            """)
+    def render(self,**kwargs):
+        super(MarkerCluster,self).render(**kwargs)
+
+        figure = self.get_root()
+        assert isinstance(figure,Figure), ("You cannot render this Element "
+            "if it's not in a Figure.")
+
+        figure.header.add_children(\
+            JavascriptLink("https://cdnjs.cloudflare.com/ajax/libs/leaflet.markercluster/0.4.0/leaflet.markercluster.js"),
+                                   name='markerclusterjs')
+        figure.header.add_children(\
+            CssLink("https://cdnjs.cloudflare.com/ajax/libs/leaflet.markercluster/0.4.0/MarkerCluster.css"),
+                                   name='markerclustercss')
+        figure.header.add_children(\
+            CssLink("https://cdnjs.cloudflare.com/ajax/libs/leaflet.markercluster/0.4.0/MarkerCluster.Default.css"),
+                                   name='markerclusterdefaultcss')

tests/test_plugins.py

@@ -28,8 +28,8 @@ def test_marker_cluster(self):
             range(N),                                    # Popups.
             ]).T
         mapa = folium.Map([45., 3.], zoom_start=4)
-        mapa.add_plugin(plugins.MarkerCluster(data))
-        mapa._build_map()
+        mapa.add_children(plugins.MarkerCluster(data))
+        mapa._repr_html_()
 
     def test_terminator(self):
         mapa = folium.Map([45., 3.], zoom_start=1)