Merge pull request #2051 from plotly/px_legend_title

nicolaskruchten · web-flow · commit 13a1c2e78ac0 · 2020-01-09T16:29:39.000-05:00
use legend titles in PX
diff --git a/doc/python/plotly-express.md b/doc/python/plotly-express.md
@@ -5,8 +5,8 @@ jupyter:
     text_representation:
       extension: .md
       format_name: markdown
-      format_version: "1.1"
-      jupytext_version: 1.1.1
+      format_version: '1.2'
+      jupytext_version: 1.3.1
   kernelspec:
     display_name: Python 3
     language: python
@@ -20,10 +20,9 @@ jupyter:
     name: python
     nbconvert_exporter: python
     pygments_lexer: ipython3
-    version: 3.7.3
+    version: 3.6.8
   plotly:
-    description:
-      Plotly Express is a terse, consistent, high-level API for rapid data
+    description: Plotly Express is a terse, consistent, high-level API for rapid data
       exploration and figure generation.
     display_as: file_settings
     language: python
diff --git a/packages/python/plotly/plotly/express/_core.py b/packages/python/plotly/plotly/express/_core.py
@@ -1184,6 +1184,7 @@ def make_figure(args, constructor, trace_patch={}, layout_patch={}):
     nrows = ncols = 1
     col_labels = []
     row_labels = []
+    trace_name_labels = None
     for group_name in sorted_group_names:
         group = grouped.get_group(group_name if len(group_name) > 1 else group_name[0])
         mapping_labels = OrderedDict()
@@ -1197,7 +1198,7 @@ def make_figure(args, constructor, trace_patch={}, layout_patch={}):
                     trace_name_labels[key] = str(val)
                 if m.variable == "animation_frame":
                     frame_name = val
-        trace_name = ", ".join(k + "=" + v for k, v in trace_name_labels.items())
+        trace_name = ", ".join(trace_name_labels.values())
         if frame_name not in trace_names_by_frame:
             trace_names_by_frame[frame_name] = set()
         trace_names = trace_names_by_frame[frame_name]
@@ -1346,7 +1347,9 @@ def make_figure(args, constructor, trace_patch={}, layout_patch={}):
     for v in ["title", "height", "width"]:
         if args[v]:
             layout_patch[v] = args[v]
-    layout_patch["legend"] = {"tracegroupgap": 0}
+    layout_patch["legend"] = dict(tracegroupgap=0)
+    if trace_name_labels:
+        layout_patch["legend"]["title"] = ", ".join(trace_name_labels)
     if "title" not in layout_patch and args["template"].layout.margin.t is None:
         layout_patch["margin"] = {"t": 60}
     if (
diff --git a/packages/python/plotly/plotly/tests/test_core/test_px/test_px.py b/packages/python/plotly/plotly/tests/test_core/test_px/test_px.py
@@ -63,18 +63,18 @@ def test_labels():
         facet_col="day",
         color="size",
         symbol="sex",
-        labels={c: c[::-1] for c in tips.columns},
+        labels={c: c.upper() for c in tips.columns},
     )
-    assert "xes" in fig.data[0].hovertemplate
-    assert "llib_latot" in fig.data[0].hovertemplate
-    assert "ezis" in fig.data[0].hovertemplate
-    assert "yad" in fig.data[0].hovertemplate
-    assert "emit" in fig.data[0].hovertemplate
-    assert fig.data[0].name.startswith("xes")
-    assert fig.layout.xaxis.title.text == "llib_latot"
-    assert fig.layout.coloraxis.colorbar.title.text == "ezis"
-    assert fig.layout.annotations[0].text.startswith("yad")
-    assert fig.layout.annotations[4].text.startswith("emit")
+    assert "SEX" in fig.data[0].hovertemplate
+    assert "TOTAL_BILL" in fig.data[0].hovertemplate
+    assert "SIZE" in fig.data[0].hovertemplate
+    assert "DAY" in fig.data[0].hovertemplate
+    assert "TIME" in fig.data[0].hovertemplate
+    assert fig.layout.legend.title.text.startswith("SEX")
+    assert fig.layout.xaxis.title.text == "TOTAL_BILL"
+    assert fig.layout.coloraxis.colorbar.title.text == "SIZE"
+    assert fig.layout.annotations[0].text.startswith("DAY")
+    assert fig.layout.annotations[4].text.startswith("TIME")
 
 
 def test_px_templates():
