continuous color scale for sunburst and treemap

Author: emmanuelle

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

@@ -1121,6 +1121,10 @@ def pie(
     data_frame=None,
     names=None,
     values=None,
+    color=None,
+    color_continuous_scale=None,
+    range_color=None,
+    color_continuous_midpoint=None,
     textinfo=None,
     hover_name=None,
     hover_data=None,
@@ -1160,6 +1164,10 @@ def sunburst(
     values=None,
     parents=None,
     ids=None,
+    color=None,
+    color_continuous_scale=None,
+    range_color=None,
+    color_continuous_midpoint=None,
     hover_name=None,
     hover_data=None,
     custom_data=None,
@@ -1191,6 +1199,10 @@ def treemap(
     values=None,
     parents=None,
     ids=None,
+    color=None,
+    color_continuous_scale=None,
+    range_color=None,
+    color_continuous_midpoint=None,
     hover_name=None,
     hover_data=None,
     custom_data=None,
@@ -1272,7 +1284,9 @@ def funnel_area(
     """
     In a funnel area plot, each row of `data_frame` is represented as a trapezoidal sector of a funnel.
     """
-    return make_figure(args=locals(), constructor=go.Funnelarea,)
+    return make_figure(
+        args=locals(), constructor=go.Funnelarea, trace_patch=dict(showlegend=True)
+    )
 
 
 funnel_area.__doc__ = make_docstring(

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

@@ -291,6 +291,13 @@ 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]:
+                    colorable = "marker"
+                    if colorable not in result:
+                        result[colorable] = dict()
+                    result[colorable]["colors"] = g[v]
+                    result[colorable]["coloraxis"] = "coloraxis1"
+                    mapping_labels[v_label] = "%{color}"
                 else:
                     colorable = "marker"
                     if trace_spec.constructor in [go.Parcats, go.Parcoords]:

packages/python/plotly/plotly/tests/test_core/test_px/test_px_functions.py

@@ -50,3 +50,31 @@ def test_pie_like_px():
     fig = px.funnel_area(values=y, names=x)
     trace = go.Funnelarea(values=y, labels=x)
     _compare_figures(trace, fig)
+
+
+def test_pie_like_colors():
+    labels = ["Eve", "Cain", "Seth", "Enos", "Noam", "Abel", "Awan", "Enoch", "Azura"]
+    parents = ["", "Eve", "Eve", "Seth", "Seth", "Eve", "Eve", "Awan", "Eve"]
+    values = [10, 14, 12, 10, 2, 6, 6, 4, 4]
+    # Sunburst
+    fig = px.sunburst(
+        names=labels,
+        parents=parents,
+        values=values,
+        color=values,
+        color_continuous_scale="Viridis",
+        range_color=(5, 15),
+    )
+    assert fig.layout.coloraxis.cmin, fig.layout.coloraxis.cmax == (5, 15)
+    assert fig.layout.coloraxis.colorscale[0] == (0.0, "#440154")
+    # Treemap
+    fig = px.treemap(
+        names=labels,
+        parents=parents,
+        values=values,
+        color=values,
+        color_continuous_scale="Viridis",
+        range_color=(5, 15),
+    )
+    assert fig.layout.coloraxis.cmin, fig.layout.coloraxis.cmax == (5, 15)
+    assert fig.layout.coloraxis.colorscale[0] == (0.0, "#440154")