imshow now handles DataFrame indexes and names by default

Author: nicolaskruchten

Diff for: doc/python/heatmaps.md

@@ -36,7 +36,7 @@ jupyter:
 
 ### Heatmap with `plotly.express` and `px.imshow`
 
-[Plotly Express](/python/plotly-express/) is the easy-to-use, high-level interface to Plotly, which [operates on a variety of types of data](/python/px-arguments/) and produces [easy-to-style figures](/python/styling-plotly-express/). With `px.imshow`, each value of the input array is represented as a heatmap pixel.
+[Plotly Express](/python/plotly-express/) is the easy-to-use, high-level interface to Plotly, which [operates on a variety of types of data](/python/px-arguments/) and produces [easy-to-style figures](/python/styling-plotly-express/). With `px.imshow`, each value of the input array or data frame is represented as a heatmap pixel.
 
 For more examples using `px.imshow`, see the [tutorial on displaying image data with plotly](/python/imshow).
 
@@ -49,6 +49,14 @@ fig = px.imshow([[1, 20, 30],
 fig.show()
 ```
 
+```python
+import plotly.express as px
+
+df = px.data.medals_wide(indexed=True)
+fig = px.imshow(df)
+fig.show()
+```
+
 ### Customizing the axes and labels on a heatmap
 
 You can use the `x`, `y` and `labels` arguments to customize the display of a heatmap, and use `.update_xaxes()` to move the x axis tick labels to the top:
@@ -182,4 +190,4 @@ Arrays of rasterized values build by datashader can be visualized using
 plotly's heatmaps, as shown in the [plotly and datashader tutorial](/python/datashader/).
 
 #### Reference
-See https://plotly.com/python/reference/#heatmap for more information and chart attribute options!
+See https://plotly.com/python/reference/#heatmap for more information and chart attribute options!

Diff for: doc/python/px-arguments.md

@@ -49,7 +49,7 @@ There are three common conventions for storing column-oriented data, usually in
 * **wide-form data** has one row per value of one of the first variable, and one column per value of the second variable. This is suitable for storing and displaying 2-dimensional data.
 * **mixed-form data** is a hybrid of long-form and wide-form data, with one row per value of one variable, and some columns representing values of another, and some columns representing more variables. See the [wide-form documentation](/python/wide-form/) for examples of how to use Plotly Express to visualize this kind of data.
 
-Every Plotly Express function other than `imshow` can operate on long-form data, and in addition, the following 2D-Cartesian functions can operate on wide-form and mixed-form data: `px.scatter`, `px.line`, `px.area`, `px.bar`, `px.histogram`, `px.violin`, `px.box`, `px.strip`, `px.funnel`, `px.density_heatmap` and `px.density_contour`.
+Every Plotly Express function can operate on long-form data (other than `px.imshow` which operates only on wide-form input), and in addition, the following 2D-Cartesian functions can operate on wide-form and mixed-form data: `px.scatter`, `px.line`, `px.area`, `px.bar`, `px.histogram`, `px.violin`, `px.box`, `px.strip`, `px.funnel`, `px.density_heatmap` and `px.density_contour`.
 
 By way of example here is the same data, represented in long-form first, and then in wide-form:
 
@@ -241,4 +241,4 @@ fig = px.bar(df, x='year', y=gdp, color='continent', labels={'y':'log gdp'},
              hover_data=['country'],
              title='Evolution of world GDP')
 fig.show()
-```
+```

Diff for: doc/python/wide-form.md

@@ -48,7 +48,7 @@ There are three common conventions for storing column-oriented data, usually in
 * **wide-form data** has one row per value of one of the first variable, and one column per value of the second variable. This is suitable for storing and displaying 2-dimensional data.
 * **mixed-form data** is a hybrid of long-form and wide-form data, with one row per value of one variable, and some columns representing values of another, and some columns representing more variables.
 
-Every Plotly Express function other than `imshow` can operate on long-form data, and in addition, the following 2D-Cartesian functions can operate on wide-form and mixed-form data: `px.scatter`, `px.line`, `px.area`, `px.bar`, `px.histogram`, `px.violin`, `px.box`, `px.strip`, `px.funnel`, `px.density_heatmap` and `px.density_contour`.
+Every Plotly Express function can operate on long-form data (other than `px.imshow` which operates only on wide-form input), and in addition, the following 2D-Cartesian functions can operate on wide-form and mixed-form data: `px.scatter`, `px.line`, `px.area`, `px.bar`, `px.histogram`, `px.violin`, `px.box`, `px.strip`, `px.funnel`, `px.density_heatmap` and `px.density_contour`.
 
 By way of example here is the same data, represented in long-form first, and then in wide-form:
 
@@ -302,4 +302,4 @@ fig.show()
 
 fig = px.box(wide_df, orientation="h")
 fig.show()
-```
+```

Diff for: packages/python/plotly/plotly/express/_imshow.py

@@ -196,6 +196,17 @@ def imshow(
             labels["color"] = xarray.plot.utils.label_from_attrs(img)
             labels["color"] = labels["color"].replace("\n", "<br>")
     else:
+        if hasattr(img, "columns") and hasattr(img.columns, "__len__"):
+            if x is None:
+                x = img.columns
+            if labels.get("x", None) is None and hasattr(img.columns, "name"):
+                labels["x"] = img.columns.name or ""
+        if hasattr(img, "index") and hasattr(img.index, "__len__"):
+            if y is None:
+                y = img.index
+            if labels.get("y", None) is None and hasattr(img.index, "name"):
+                labels["y"] = img.index.name or ""
+
         if labels.get("x", None) is None:
             labels["x"] = ""
         if labels.get("y", None) is None:

Diff for: packages/python/plotly/plotly/tests/test_core/test_px/test_imshow.py

@@ -150,3 +150,25 @@ def test_imshow_labels_and_ranges():
 
     with pytest.raises(ValueError):
         fig = px.imshow([[1, 2], [3, 4], [5, 6]], x=["a"])
+
+
+def test_imshow_dataframe():
+    df = px.data.medals_wide(indexed=False)
+    fig = px.imshow(df)
+    assert fig.data[0].x[0] == df.columns[0]
+    assert fig.data[0].x[0] == "nation"
+    assert fig.layout.xaxis.title.text is None
+    assert fig.data[0].y[0] == df.index[0]
+    assert fig.data[0].y[0] == 0
+    assert fig.layout.yaxis.title.text is None
+
+    df = px.data.medals_wide(indexed=True)
+    fig = px.imshow(df)
+    assert fig.data[0].x[0] == df.columns[0]
+    assert fig.data[0].x[0] == "gold"
+    assert fig.layout.xaxis.title.text == df.columns.name
+    assert fig.layout.xaxis.title.text == "medal"
+    assert fig.data[0].y[0] == df.index[0]
+    assert fig.data[0].y[0] == "South Korea"
+    assert fig.layout.yaxis.title.text == df.index.name
+    assert fig.layout.yaxis.title.text == "nation"