Add style_function to TopoJson

Author: Martin Journois

folium/features.py

@@ -308,7 +308,7 @@ def _get_self_bounds(self):
         return bounds
 
 class TopoJson(MacroElement):
-    def __init__(self, data, object_path):
+    def __init__(self, data, object_path, style_function=None):
         """
         TODO docstring here
 
@@ -317,26 +317,44 @@ def __init__(self, data, object_path):
         self._name = 'TopoJson'
         if 'read' in dir(data):
             self.embed = True
-            self.data = data.read()
+            self.data = json.load(data)
         elif type(data) is dict:
             self.embed = True
-            self.data = json.dumps(data)
+            self.data = data
         else:
             self.embed = False
             self.data = data
 
         self.object_path = object_path
 
+        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()}}_data = {{this.data}};
+                var {{this.get_name()}}_data = {{this.style_data()}};
                 var {{this.get_name()}} = L.geoJson(topojson.feature(
                     {{this.get_name()}}_data,
                     {{this.get_name()}}_data.{{this.object_path}}
                     )).addTo({{this._parent.get_name()}});
+                {{this.get_name()}}.setStyle(function(feature) {return feature.properties.style;});
+
             {% endmacro %}
             """)
 
+    def style_data(self):
+        def recursive_get(data, keys):
+            if len(keys):
+                return recursive_get(data.get(keys[0]), keys[1:])
+            else:
+                return data
+        geometries = recursive_get(self.data, self.object_path.split('.'))['geometries']
+        for feature in geometries:
+            feature.setdefault('properties',{}).setdefault('style',{}).update(
+                self.style_function(feature))
+        return json.dumps(self.data)
+
     def render(self, **kwargs):
         super(TopoJson, self).render(**kwargs)
 

folium/folium.py

@@ -610,7 +610,7 @@ def choropleth(self, geo_path=None, geo_str=None, data_out='data.json',
             }
 
         if topojson:
-            geo_json = TopoJson(geo_data, topojson)
+            geo_json = TopoJson(geo_data, topojson, style_function=style_function)
         else:
             geo_json = GeoJson(geo_data, style_function=style_function)