Merge pull request #1909 from plotly/px-pie-etc

adding pie, sunburst, funnel and funnelarea to px

packages/python/plotly/plotly/express/__init__.py

@@ -39,6 +39,11 @@
     choropleth,
     density_contour,
     density_heatmap,
+    pie,
+    sunburst,
+    treemap,
+    funnel,
+    funnel_area,
 )
 
 from ._imshow import imshow
@@ -77,6 +82,11 @@
     "strip",
     "histogram",
     "choropleth",
+    "pie",
+    "sunburst",
+    "treemap",
+    "funnel",
+    "funnel_area",
     "imshow",
     "data",
     "colors",

packages/python/plotly/plotly/express/_chart_types.py

@@ -1115,3 +1115,208 @@ def parallel_categories(
 
 
 parallel_categories.__doc__ = make_docstring(parallel_categories)
+
+
+def pie(
+    data_frame=None,
+    names=None,
+    values=None,
+    color=None,
+    color_discrete_sequence=None,
+    color_discrete_map={},
+    hover_name=None,
+    hover_data=None,
+    custom_data=None,
+    labels={},
+    title=None,
+    template=None,
+    width=None,
+    height=None,
+    opacity=None,
+    hole=None,
+):
+    """
+    In a pie plot, each row of `data_frame` is represented as a sector of a pie.
+    """
+    if color_discrete_sequence is not None:
+        layout_patch = {"piecolorway": color_discrete_sequence}
+    else:
+        layout_patch = {}
+    return make_figure(
+        args=locals(),
+        constructor=go.Pie,
+        trace_patch=dict(showlegend=(names is not None), hole=hole),
+        layout_patch=layout_patch,
+    )
+
+
+pie.__doc__ = make_docstring(
+    pie,
+    override_dict=dict(
+        hole=[
+            "float",
+            "Sets the fraction of the radius to cut out of the pie."
+            "Use this to make a donut chart.",
+        ],
+    ),
+)
+
+
+def sunburst(
+    data_frame=None,
+    names=None,
+    values=None,
+    parents=None,
+    ids=None,
+    color=None,
+    color_continuous_scale=None,
+    range_color=None,
+    color_continuous_midpoint=None,
+    color_discrete_sequence=None,
+    color_discrete_map={},
+    hover_name=None,
+    hover_data=None,
+    custom_data=None,
+    labels={},
+    title=None,
+    template=None,
+    width=None,
+    height=None,
+    branchvalues=None,
+    maxdepth=None,
+):
+    """
+    A sunburst plot represents hierarchial data as sectors laid out over
+    several levels of concentric rings.
+    """
+    if color_discrete_sequence is not None:
+        layout_patch = {"sunburstcolorway": color_discrete_sequence}
+    else:
+        layout_patch = {}
+    return make_figure(
+        args=locals(),
+        constructor=go.Sunburst,
+        trace_patch=dict(branchvalues=branchvalues, maxdepth=maxdepth),
+        layout_patch=layout_patch,
+    )
+
+
+sunburst.__doc__ = make_docstring(sunburst)
+
+
+def treemap(
+    data_frame=None,
+    names=None,
+    values=None,
+    parents=None,
+    ids=None,
+    color=None,
+    color_continuous_scale=None,
+    range_color=None,
+    color_continuous_midpoint=None,
+    color_discrete_sequence=None,
+    color_discrete_map={},
+    hover_name=None,
+    hover_data=None,
+    custom_data=None,
+    labels={},
+    title=None,
+    template=None,
+    width=None,
+    height=None,
+    branchvalues=None,
+    maxdepth=None,
+):
+    """
+    A treemap plot represents hierarchial data as nested rectangular sectors.
+    """
+    if color_discrete_sequence is not None:
+        layout_patch = {"treemapcolorway": color_discrete_sequence}
+    else:
+        layout_patch = {}
+    return make_figure(
+        args=locals(),
+        constructor=go.Treemap,
+        trace_patch=dict(branchvalues=branchvalues, maxdepth=maxdepth),
+        layout_patch=layout_patch,
+    )
+
+
+treemap.__doc__ = make_docstring(treemap)
+
+
+def funnel(
+    data_frame=None,
+    x=None,
+    y=None,
+    color=None,
+    facet_row=None,
+    facet_col=None,
+    facet_col_wrap=0,
+    hover_name=None,
+    hover_data=None,
+    custom_data=None,
+    text=None,
+    animation_frame=None,
+    animation_group=None,
+    category_orders={},
+    labels={},
+    color_discrete_sequence=None,
+    color_discrete_map={},
+    opacity=None,
+    orientation="h",
+    log_x=False,
+    log_y=False,
+    range_x=None,
+    range_y=None,
+    title=None,
+    template=None,
+    width=None,
+    height=None,
+):
+    """
+    In a funnel plot, each row of `data_frame` is represented as a rectangular sector of a funnel.
+    """
+    return make_figure(
+        args=locals(),
+        constructor=go.Funnel,
+        trace_patch=dict(opacity=opacity, orientation=orientation),
+    )
+
+
+funnel.__doc__ = make_docstring(funnel)
+
+
+def funnel_area(
+    data_frame=None,
+    names=None,
+    values=None,
+    color=None,
+    color_discrete_sequence=None,
+    color_discrete_map={},
+    hover_name=None,
+    hover_data=None,
+    custom_data=None,
+    labels={},
+    title=None,
+    template=None,
+    width=None,
+    height=None,
+    opacity=None,
+):
+    """
+    In a funnel area plot, each row of `data_frame` is represented as a trapezoidal sector of a funnel.
+    """
+    if color_discrete_sequence is not None:
+        layout_patch = {"funnelareacolorway": color_discrete_sequence}
+    else:
+        layout_patch = {}
+    return make_figure(
+        args=locals(),
+        constructor=go.Funnelarea,
+        trace_patch=dict(showlegend=(names is not None)),
+        layout_patch=layout_patch,
+    )
+
+
+funnel_area.__doc__ = make_docstring(funnel_area)

packages/python/plotly/plotly/express/_core.py

@@ -291,6 +291,28 @@ def make_trace_kwargs(args, trace_spec, g, mapping_labels, sizeref):
                     result["z"] = g[v]
                     result["coloraxis"] = "coloraxis1"
                     mapping_labels[v_label] = "%{z}"
+                elif trace_spec.constructor in [
+                    go.Sunburst,
+                    go.Treemap,
+                    go.Pie,
+                    go.Funnelarea,
+                ]:
+                    if "marker" not in result:
+                        result["marker"] = dict()
+
+                    if args.get("color_is_continuous"):
+                        result["marker"]["colors"] = g[v]
+                        result["marker"]["coloraxis"] = "coloraxis1"
+                        mapping_labels[v_label] = "%{color}"
+                    else:
+                        result["marker"]["colors"] = []
+                        mapping = {}
+                        for cat in g[v]:
+                            if mapping.get(cat) is None:
+                                mapping[cat] = args["color_discrete_sequence"][
+                                    len(mapping) % len(args["color_discrete_sequence"])
+                                ]
+                            result["marker"]["colors"].append(mapping[cat])
                 else:
                     colorable = "marker"
                     if trace_spec.constructor in [go.Parcats, go.Parcoords]:
@@ -305,11 +327,38 @@ def make_trace_kwargs(args, trace_spec, g, mapping_labels, sizeref):
             elif k == "locations":
                 result[k] = g[v]
                 mapping_labels[v_label] = "%{location}"
+            elif k == "values":
+                result[k] = g[v]
+                _label = "value" if v_label == "values" else v_label
+                mapping_labels[_label] = "%{value}"
+            elif k == "parents":
+                result[k] = g[v]
+                _label = "parent" if v_label == "parents" else v_label
+                mapping_labels[_label] = "%{parent}"
+            elif k == "ids":
+                result[k] = g[v]
+                _label = "id" if v_label == "ids" else v_label
+                mapping_labels[_label] = "%{id}"
+            elif k == "names":
+                if trace_spec.constructor in [
+                    go.Sunburst,
+                    go.Treemap,
+                    go.Pie,
+                    go.Funnelarea,
+                ]:
+                    result["labels"] = g[v]
+                    _label = "label" if v_label == "names" else v_label
+                    mapping_labels[_label] = "%{label}"
+                else:
+                    result[k] = g[v]
             else:
                 if v:
                     result[k] = g[v]
                 mapping_labels[v_label] = "%%{%s}" % k
-    if trace_spec.constructor not in [go.Parcoords, go.Parcats]:
+    if trace_spec.constructor not in [
+        go.Parcoords,
+        go.Parcats,
+    ]:
         hover_lines = [k + "=" + v for k, v in mapping_labels.items()]
         result["hovertemplate"] = hover_header + "<br>".join(hover_lines)
     return result, fit_results
@@ -674,6 +723,7 @@ def one_group(x):
 
 def apply_default_cascade(args):
     # first we apply px.defaults to unspecified args
+
     for param in (
         ["color_discrete_sequence", "color_continuous_scale"]
         + ["symbol_sequence", "line_dash_sequence", "template"]
@@ -956,6 +1006,7 @@ def infer_config(args, constructor, trace_patch):
     attrables = (
         ["x", "y", "z", "a", "b", "c", "r", "theta", "size", "dimensions"]
         + ["custom_data", "hover_name", "hover_data", "text"]
+        + ["names", "values", "parents", "ids"]
         + ["error_x", "error_x_minus"]
         + ["error_y", "error_y_minus", "error_z", "error_z_minus"]
         + ["lat", "lon", "locations", "animation_group"]
@@ -989,14 +1040,34 @@ def infer_config(args, constructor, trace_patch):
                     and args["data_frame"][args["color"]].dtype.kind in "bifc"
                 ):
                     attrs.append("color")
+                    args["color_is_continuous"] = True
+                elif constructor in [go.Sunburst, go.Treemap]:
+                    attrs.append("color")
+                    args["color_is_continuous"] = False
                 else:
                     grouped_attrs.append("marker.color")
         elif "line_group" in args or constructor == go.Histogram2dContour:
             grouped_attrs.append("line.color")
+        elif constructor in [go.Pie, go.Funnelarea]:
+            attrs.append("color")
+            if args["color"]:
+                if args["hover_data"] is None:
+                    args["hover_data"] = []
+                args["hover_data"].append(args["color"])
         else:
             grouped_attrs.append("marker.color")
 
-        show_colorbar = bool("color" in attrs and args["color"])
+        show_colorbar = bool(
+            "color" in attrs
+            and args["color"]
+            and constructor not in [go.Pie, go.Funnelarea]
+            and (
+                constructor not in [go.Treemap, go.Sunburst]
+                or args.get("color_is_continuous")
+            )
+        )
+    else:
+        show_colorbar = False
 
     # Compute line_dash grouping attribute
     if "line_dash" in args:
@@ -1148,6 +1219,8 @@ def make_figure(args, constructor, trace_patch={}, layout_patch={}):
                 go.Parcoords,
                 go.Choropleth,
                 go.Histogram2d,
+                go.Sunburst,
+                go.Treemap,
             ]:
                 trace.update(
                     legendgroup=trace_name,