Merge pull request #266 from BibMartin/choropleth

Replace GeoJsonStyle by kwarg style_function in GeoJson

Author: themiurgo

.gitignore

@@ -44,3 +44,7 @@ nosetests.xml
 #Virtualenv
 ENV
 .env
+
+# Tests products
+.cache
+

folium/__init__.py

@@ -10,7 +10,7 @@
                         Popup, TileLayer)
 
 from folium.features import (ClickForMarker, ColorScale, CustomIcon, DivIcon,
-                             GeoJson, GeoJsonStyle, ImageOverlay, LatLngPopup,
+                             GeoJson, ImageOverlay, LatLngPopup,
                              MarkerCluster, MultiPolyLine, PolyLine, Vega,
                              RegularPolygonMarker, TopoJson, WmsTileLayer)
 

folium/features.py

@@ -194,20 +194,22 @@ def render(self, **kwargs):
 
 
 class GeoJson(MacroElement):
-    def __init__(self, data):
+    def __init__(self, data, style_function=None):
         """
         Creates a GeoJson plugin to append into a map with
         Map.add_plugin.
 
         Parameters
         ----------
         data: file, dict or str.
-        The GeoJSON data you want to plot.
-        * If file, then data will be read in the file and fully
-          embedded in Leaflet's JavaScript.
-        * If dict, then data will be converted to JSON and embedded
-          in the JavaScript.
-        * If str, then data will be passed to the JavaScript as-is.
+            The GeoJSON data you want to plot.
+            * If file, then data will be read in the file and fully
+              embedded in Leaflet's JavaScript.
+            * If dict, then data will be converted to JSON and embedded
+              in the JavaScript.
+            * If str, then data will be passed to the JavaScript as-is.
+        style_function: function, default None
+            A function mapping a GeoJson Feature to a style dict.
 
         Examples
         --------
@@ -218,22 +220,43 @@ def __init__(self, data):
         >>> # Providing string.
         >>> GeoJson(open('foo.json').read())
 
+        >>> # Providing a style_function that put all states in green, but Alabama in blue.
+        >>> style_function=lambda x: {'fillColor': '#0000ff' if x['properties']['name']=='Alabama' else '#00ff00'}
+        >>> GeoJson(geojson, style_function=style_function)
         """
         super(GeoJson, self).__init__()
         self._name = 'GeoJson'
         if 'read' in dir(data):
-            self.data = data.read()
+            self.data = json.load(data)
         elif type(data) is dict:
-            self.data = json.dumps(data)
-        else:
             self.data = data
+        else:
+            self.data = json.loads(data)
+
+        if 'features' not in self.data.keys():
+            # Catch case when GeoJSON is just a single Feature or a geometry.
+            if not (isinstance(self.data, dict) and 'geometry' in self.data.keys()):
+                # Catch case when GeoJSON is just a geometry.
+                self.data = {'type' : 'Feature', 'geometry' : self.data}
+            self.data = {'type' : 'FeatureCollection', 'features' : [self.data]}
+
+        if style_function is None:
+            style_function = lambda x: {}
+        self.style_function = style_function
 
         self._template = Template(u"""
             {% macro script(this, kwargs) %}
-                var {{this.get_name()}} = L.geoJson({{this.data}}).addTo({{this._parent.get_name()}});
+                var {{this.get_name()}} = L.geoJson({{this.style_data()}})
+                    .addTo({{this._parent.get_name()}})
+                    .setStyle(function(feature) {return feature.properties.style;});
             {% endmacro %}
             """)  # noqa
 
+    def style_data(self):
+        for feature in self.data['features']:
+            feature.setdefault('properties',{}).setdefault('style',{}).update(
+                self.style_function(feature))
+        return json.dumps(self.data)
 
 class TopoJson(MacroElement):
     def __init__(self, data, object_path):
@@ -273,81 +296,6 @@ def render(self, **kwargs):
             JavascriptLink("https://cdnjs.cloudflare.com/ajax/libs/topojson/1.6.9/topojson.min.js"),  # noqa
             name='topojson')
 
-
-class GeoJsonStyle(MacroElement):
-    def __init__(self, color_domain, color_code, color_data=None,
-                 key_on='feature.properties.color',
-                 weight=1, opacity=1, color='black',
-                 fill_opacity=0.6, dash_array=0):
-        """
-        TODO docstring here
-
-        """
-        super(GeoJsonStyle, self).__init__()
-        self._name = 'GeoJsonStyle'
-
-        self.color_domain = color_domain
-        self.color_range = color_brewer(color_code, n=len(color_domain))
-        self.color_data = json.dumps(color_data)
-        self.key_on = key_on
-
-        self.weight = weight
-        self.opacity = opacity
-        self.color = color
-        self.fill_color = color_code
-        self.fill_opacity = fill_opacity
-        self.dash_array = dash_array
-
-        self._template = Template(u"""
-            {% macro script(this, kwargs) %}
-                {% if not this.color_range %}
-                    var {{this.get_name()}} = {
-                        color_function : function(feature) {
-                            return '{{this.fill_color}}';
-                            },
-                        };
-                {%else%}
-                    var {{this.get_name()}} = {
-                        color_scale : d3.scale.threshold()
-                              .domain({{this.color_domain}})
-                              .range({{this.color_range}}),
-                        color_data : {{this.color_data}},
-                        color_function : function(feature) {
-                            {% if this.color_data=='null' %}
-                                return this.color_scale({{this.key_on}});
-                            {% else %}
-                                return
-                                  this.color_scale(this.color_data[{{this.key_on}}]);
-                            {% endif %}
-                            },
-                        };
-                {%endif%}
-
-                {{this._parent.get_name()}}.setStyle(function(feature) {
-                    return {
-                        fillColor: {{this.get_name()}}.color_function(feature),
-                        weight: {{this.weight}},
-                        opacity: {{this.opacity}},
-                        color: '{{this.color}}',
-                        fillOpacity: {{this.fill_opacity}},
-                        dashArray: '{{this.dash_array}}'
-                        };
-                    });
-            {% endmacro %}
-            """)
-
-    def render(self, **kwargs):
-        super(GeoJsonStyle, 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/d3/3.5.5/d3.min.js"),  # noqa
-            name='d3')
-
-
 class ColorScale(MacroElement):
     def __init__(self, color_domain, color_code, caption=""):
         """
@@ -505,7 +453,6 @@ def __init__(self, location, radius=500, color='black',
             {% endmacro %}
             """)
 
-
 class LatLngPopup(MacroElement):
     def __init__(self):
         """

folium/folium.py

@@ -16,7 +16,7 @@
 from .map import Map as _Map
 from .map import Icon, Marker, Popup, FitBounds
 from .features import (WmsTileLayer, RegularPolygonMarker, Vega, GeoJson,
-                       GeoJsonStyle, CircleMarker, LatLngPopup,
+                       CircleMarker, LatLngPopup,
                        ClickForMarker, ColorScale, TopoJson, PolyLine,
                        MultiPolyLine, ImageOverlay)
 from .utilities import color_brewer, write_png
@@ -442,10 +442,17 @@ def add_plugin(self, plugin):
 #
 #        self.template_vars.update({'fit_bounds': fit_bounds_str.strip()})
 
-    def geo_json(self, geo_path=None, geo_str=None, data_out='data.json',
+    def geo_json(self, *args, **kwargs):
+        """This method is deprecated and will be removed in v0.2.1. See
+        `Map.choropleth` instead.
+        """
+        warnings.warn('This method is deprecated. Please use Map.choropleth instead.')
+        return self.choropleth(*args, **kwargs)
+
+    def choropleth(self, geo_path=None, geo_str=None, data_out='data.json',
                  data=None, columns=None, key_on=None, threshold_scale=None,
                  fill_color='blue', fill_opacity=0.6, line_color='black',
-                 line_weight=1, line_opacity=1, legend_name=None,
+                 line_weight=1, line_opacity=1, legend_name="",
                  topojson=None, reset=False):
         """Apply a GeoJSON overlay to the map.
 
@@ -506,8 +513,8 @@ def geo_json(self, geo_path=None, geo_str=None, data_out='data.json',
             GeoJSON geopath line weight.
         line_opacity: float, default 1
             GeoJSON geopath line opacity, range 0-1.
-        legend_name: string, default None
-            Title for data legend. If not passed, defaults to columns[1].
+        legend_name: string, default empty string
+            Title for data legend.
         topojson: string, default None
             If using a TopoJSON, passing "objects.yourfeature" to the topojson
             keyword argument will enable conversion to GeoJSON.
@@ -520,17 +527,14 @@ def geo_json(self, geo_path=None, geo_str=None, data_out='data.json',
 
         Example
         -------
-        >>> m.geo_json(geo_path='us-states.json', line_color='blue',
+        >>> m.choropleth(geo_path='us-states.json', line_color='blue',
                       line_weight=3)
-        >>> m.geo_json(geo_path='geo.json', data=df,
+        >>> m.choropleth(geo_path='geo.json', data=df,
                       columns=['Data 1', 'Data 2'],
                       key_on='feature.properties.myvalue', fill_color='PuBu',
                       threshold_scale=[0, 20, 30, 40, 50, 60])
-        >>> m.geo_json(geo_path='countries.json', topojson='objects.countries')
+        >>> m.choropleth(geo_path='countries.json', topojson='objects.countries')
         """
-        warnings.warn("%s is deprecated. Use %s instead" %
-                      ("geo_json", "add_children(GeoJson)"),
-                      FutureWarning, stacklevel=2)
 
         if threshold_scale and len(threshold_scale) > 6:
             raise ValueError
@@ -546,11 +550,6 @@ def geo_json(self, geo_path=None, geo_str=None, data_out='data.json',
         else:
             geo_data = {}
 
-        if topojson:
-            geo_json = TopoJson(geo_data, topojson)
-        else:
-            geo_json = GeoJson(geo_data)
-
         # Create color_data dict
         if hasattr(data, 'set_index'):
             # This is a pd.DataFrame
@@ -587,20 +586,38 @@ def geo_json(self, geo_path=None, geo_str=None, data_out='data.json',
             color_domain = [data_min+i*(data_max-data_min)*1./nb_class
                             for i in range(1+nb_class)]
         else:
-            color_domain = [-1, 1]
+            color_domain = None
+
+        if color_domain and key_on:
+            key_on = key_on[8:] if key_on.startswith('feature.') else key_on
+            color_range = color_brewer(fill_color, n=len(color_domain))
+            get_by_key = lambda obj,key : (obj.get(key,None) if len(key.split('.'))<=1
+                                           else get_by_key(obj.get(key.split('.')[0],None),
+                                                           '.'.join(key.split('.')[1:])))
+            color_scale_fun = lambda x: color_range[len([u for u in color_domain
+                                                         if u<=color_data[get_by_key(x,key_on)]])]
+        else:
+            color_scale_fun = lambda x: fill_color
 
-        # Create GeoJsonStyle.
-        geo_json_style = GeoJsonStyle(
-            color_domain, fill_color, color_data=color_data, key_on=key_on,
-            weight=line_weight, opacity=line_opacity, color=line_color,
-            fill_opacity=fill_opacity)
+        style_function = lambda x: {
+            "weight" : line_weight,
+            "opactiy": line_opacity,
+            "color" : line_color,
+            "fillOpacity" : fill_opacity,
+            "fillColor" : color_scale_fun(x)
+            }
 
-        # Create ColorScale
-        color_scale = ColorScale(color_domain, fill_color, caption=legend_name)
+        if topojson:
+            geo_json = TopoJson(geo_data, topojson)
+        else:
+            geo_json = GeoJson(geo_data, style_function=style_function)
 
-        geo_json.add_children(geo_json_style)
         self.add_children(geo_json)
-        self.add_children(color_scale)
+
+        # Create ColorScale.
+        if color_domain:
+            color_scale = ColorScale(color_domain, fill_color, caption=legend_name)
+            self.add_children(color_scale)
 
     def image_overlay(self, data, opacity=0.25, min_lat=-90.0, max_lat=90.0,
                       min_lon=-180.0, max_lon=180.0, origin='upper',

tests/.gitignore

@@ -0,0 +1,3 @@
+# Tests products
+data.png
+map.html