add choropleth_mapbox to px

nicolaskruchten · nicolaskruchten · commit dd577a0a7b16 · 2019-11-30T20:46:08.000-05:00
diff --git a/packages/python/plotly/plotly/express/__init__.py b/packages/python/plotly/plotly/express/__init__.py
@@ -44,6 +44,7 @@
     treemap,
     funnel,
     funnel_area,
+    choropleth_mapbox,
 )
 
 from ._imshow import imshow
@@ -82,6 +83,7 @@
     "strip",
     "histogram",
     "choropleth",
+    "choropleth_mapbox",
     "pie",
     "sunburst",
     "treemap",
diff --git a/packages/python/plotly/plotly/express/_chart_types.py b/packages/python/plotly/plotly/express/_chart_types.py
@@ -842,15 +842,13 @@ def choropleth(
     hover_name=None,
     hover_data=None,
     custom_data=None,
-    size=None,
     animation_frame=None,
     animation_group=None,
     category_orders={},
     labels={},
     color_continuous_scale=None,
     range_color=None,
     color_continuous_midpoint=None,
-    size_max=None,
     projection=None,
     scope=None,
     center=None,
@@ -998,6 +996,42 @@ def scatter_mapbox(
 scatter_mapbox.__doc__ = make_docstring(scatter_mapbox)
 
 
+def choropleth_mapbox(
+    data_frame=None,
+    geojson=None,
+    locations=None,
+    color=None,
+    hover_name=None,
+    hover_data=None,
+    custom_data=None,
+    animation_frame=None,
+    animation_group=None,
+    category_orders={},
+    labels={},
+    color_continuous_scale=None,
+    range_color=None,
+    color_continuous_midpoint=None,
+    opacity=None,
+    zoom=None,
+    title=None,
+    template=None,
+    width=None,
+    height=None,
+):
+    """
+    In a Mapbox choropleth map, each row of `data_frame` is represented by a
+    colored region on a Mapbox map.
+    """
+    return make_figure(
+        args=locals(),
+        constructor=go.Choroplethmapbox,
+        trace_patch=dict(geojson=geojson),
+    )
+
+
+choropleth_mapbox.__doc__ = make_docstring(choropleth_mapbox)
+
+
 def line_mapbox(
     data_frame=None,
     lat=None,
diff --git a/packages/python/plotly/plotly/express/_core.py b/packages/python/plotly/plotly/express/_core.py
@@ -287,7 +287,7 @@ def make_trace_kwargs(args, trace_spec, g, mapping_labels, sizeref):
                         v_label_col = get_decorated_label(args, col, None)
                         mapping_labels[v_label_col] = "%%{customdata[%d]}" % (position)
             elif k == "color":
-                if trace_spec.constructor == go.Choropleth:
+                if trace_spec.constructor in [go.Choropleth, go.Choroplethmapbox]:
                     result["z"] = g[v]
                     result["coloraxis"] = "coloraxis1"
                     mapping_labels[v_label] = "%{z}"
@@ -380,6 +380,7 @@ def configure_axes(args, constructor, fig, orders):
         go.Scatterpolargl: configure_polar_axes,
         go.Barpolar: configure_polar_axes,
         go.Scattermapbox: configure_mapbox,
+        go.Choroplethmapbox: configure_mapbox,
         go.Scattergeo: configure_geo,
         go.Choropleth: configure_geo,
     }
@@ -569,7 +570,9 @@ def configure_mapbox(args, fig, orders):
                 center=dict(
                     lat=args["data_frame"][args["lat"]].mean(),
                     lon=args["data_frame"][args["lon"]].mean(),
-                ),
+                )
+                if "lat" in args and "lon" in  args
+                else dict(),
                 zoom=args["zoom"],
             )
         )
@@ -1221,6 +1224,7 @@ def make_figure(args, constructor, trace_patch={}, layout_patch={}):
                 go.Parcats,
                 go.Parcoords,
                 go.Choropleth,
+                go.Choroplethmapbox,
                 go.Histogram2d,
                 go.Sunburst,
                 go.Treemap,
diff --git a/packages/python/plotly/plotly/express/_doc.py b/packages/python/plotly/plotly/express/_doc.py
@@ -456,6 +456,10 @@
     ],
     box=["boolean (default `False`)", "If `True`, boxes are drawn inside the violins."],
     notched=["boolean (default `False`)", "If `True`, boxes are drawn with notches."],
+    geojson=[
+        "GeoJSON-formatted dict",
+        "Must contain a Polygon feature collection, with IDs, which are references from `locations`.",
+    ],
     cumulative=[
         "boolean (default `False`)",
         "If `True`, histogram values are cumulative.",
