Use a dataclass instead of a dict as a return type for boxplot()

timhoffm · timhoffm · commit ef3d7e975e89 · 2024-02-16T12:19:47.000+01:00
This is a prove of concept. Backward compatibility is achieved though
a collections.abc.Mapping mixin. Note though, that I did not go as far
as MutableMapping. If someone has written to the returned dict, they
would still be broken.

open issues:
- documentation incomplete
- should switch usages in the matplotlib code to attributes
- Where to put the dataclass? cbook or axes?
- should we use silent_list to make the repr a bit nicer?
diff --git a/doc/api/next_api_changes/behavior/27788-TH.rst b/doc/api/next_api_changes/behavior/27788-TH.rst
@@ -0,0 +1,3 @@
+``Boxplot()`` returns a dataclass instead of a dict
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The dataclass still supports dict-like access for backwards compatibility.
diff --git a/galleries/examples/statistics/boxplot_color.py b/galleries/examples/statistics/boxplot_color.py
@@ -29,7 +29,7 @@
                    labels=labels)  # will be used to label x-ticks
 
 # fill with colors
-for patch, color in zip(bplot['boxes'], colors):
+for patch, color in zip(bplot.boxes, colors):
     patch.set_facecolor(color)
 
 plt.show()
diff --git a/galleries/examples/statistics/boxplot_demo.py b/galleries/examples/statistics/boxplot_demo.py
@@ -108,9 +108,9 @@
 fig.subplots_adjust(left=0.075, right=0.95, top=0.9, bottom=0.25)
 
 bp = ax1.boxplot(data, notch=False, sym='+', vert=True, whis=1.5)
-plt.setp(bp['boxes'], color='black')
-plt.setp(bp['whiskers'], color='black')
-plt.setp(bp['fliers'], color='red', marker='+')
+plt.setp(bp.boxes, color='black')
+plt.setp(bp.whiskers, color='black')
+plt.setp(bp.fliers, color='red', marker='+')
 
 # Add a horizontal grid to the plot, but make it very light in color
 # so we can use it for reading data values but not be distracting
@@ -229,8 +229,8 @@ def fake_bootstrapper(n):
 
 ax.set_xlabel('treatment')
 ax.set_ylabel('response')
-plt.setp(bp['whiskers'], color='k', linestyle='-')
-plt.setp(bp['fliers'], markersize=3.0)
+plt.setp(bp.whiskers, color='k', linestyle='-')
+plt.setp(bp.fliers, markersize=3.0)
 plt.show()
 
 
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3908,10 +3908,10 @@ def boxplot(self, x, notch=None, sym=None, vert=None, whis=None,
 
         Returns
         -------
-        dict
-          A dictionary mapping each component of the boxplot to a list
-          of the `.Line2D` instances created. That dictionary has the
-          following keys (assuming vertical boxplots):
+        BoxplotArtists
+          A dataclass mapping each component of the boxplot to a list
+          of the `.Line2D` instances created. The dataclass has the
+          following attributes (assuming vertical boxplots):
 
           - ``boxes``: the main body of the boxplot showing the
             quartiles and the median's confidence intervals if
@@ -3930,6 +3930,10 @@ def boxplot(self, x, notch=None, sym=None, vert=None, whis=None,
 
           - ``means``: points or lines representing the means.
 
+          .. versionchanged:: 3.9
+              Formerly, a dict was returned. The return value is now a dataclass, that
+              still supports dict-like access for backward-compatibility.
+
         Other Parameters
         ----------------
         showcaps : bool, default: True
@@ -4157,10 +4161,10 @@ def bxp(self, bxpstats, positions=None, widths=None, vert=True,
 
         Returns
         -------
-        dict
-          A dictionary mapping each component of the boxplot to a list
-          of the `.Line2D` instances created. That dictionary has the
-          following keys (assuming vertical boxplots):
+        BoxplotArtists
+          A dataclass mapping each component of the boxplot to a list
+          of the `.Line2D` instances created. The dataclass has the
+          following attributes (assuming vertical boxplots):
 
           - ``boxes``: main bodies of the boxplot showing the quartiles, and
             the median's confidence intervals if enabled.
@@ -4170,6 +4174,10 @@ def bxp(self, bxpstats, positions=None, widths=None, vert=True,
           - ``fliers``: points representing data beyond the whiskers (fliers).
           - ``means``: points or lines representing the means.
 
+        .. versionchanged:: 3.9
+            Formerly, a dict was returned. The return value is now a dataclass, that
+            still supports dict-like access for backward-compatibility.
+
         Examples
         --------
         .. plot:: gallery/statistics/bxp.py
@@ -4376,8 +4384,10 @@ def do_patch(xs, ys, **kwargs):
 
             self._request_autoscale_view()
 
-        return dict(whiskers=whiskers, caps=caps, boxes=boxes,
-                    medians=medians, fliers=fliers, means=means)
+        return cbook.BoxplotArtists(
+            whiskers=whiskers, caps=caps, boxes=boxes,
+            medians=medians, fliers=fliers, means=means
+        )
 
     @staticmethod
     def _parse_scatter_color_args(c, edgecolors, kwargs, xsize,
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -6,6 +6,7 @@
 import collections
 import collections.abc
 import contextlib
+import dataclasses
 import functools
 import gzip
 import itertools
@@ -19,6 +20,7 @@
 import time
 import traceback
 import types
+from typing import Union
 import weakref
 
 import numpy as np
@@ -1126,6 +1128,31 @@ def _broadcast_with_masks(*args, compress=False):
     return inputs
 
 
+@dataclasses.dataclass(frozen=True)
+class BoxplotArtists(collections.abc.Mapping):
+    """
+    Collection of the artists created by `~.Axes.boxplot`.
+
+    For backward compatibility with the previously used dict representation,
+    this can alternatively be accessed like a (read-only) dict.
+    """
+    boxes: Union[list["matplotlib.lines.Line2D"], list["matplotlib.patches.Patch"]]
+    medians: list["matplotlib.lines.Line2D"]
+    whiskers: list["matplotlib.lines.Line2D"]
+    caps: list["matplotlib.lines.Line2D"]
+    fliers: list["matplotlib.lines.Line2D"]
+    means: list["matplotlib.lines.Line2D"]
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __iter__(self):
+        return iter(self.__annotations__)
+
+    def __len__(self):
+        return len(self.__annotations__)
+
+
 def boxplot_stats(X, whis=1.5, bootstrap=None, labels=None,
                   autorange=False):
     r"""
diff --git a/lib/matplotlib/cbook.pyi b/lib/matplotlib/cbook.pyi
@@ -1,10 +1,13 @@
 import collections.abc
 from collections.abc import Callable, Collection, Generator, Iterable, Iterator
 import contextlib
+import dataclasses
 import os
 from pathlib import Path
 
 from matplotlib.artist import Artist
+from matplotlib.lines import Line2D
+from matplotlib.patches import Patch
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -15,6 +18,7 @@ from typing import (
     IO,
     Literal,
     TypeVar,
+    Union,
     overload,
 )
 
@@ -131,6 +135,19 @@ class GrouperView(Generic[_T]):
 def simple_linear_interpolation(a: ArrayLike, steps: int) -> np.ndarray: ...
 def delete_masked_points(*args): ...
 def _broadcast_with_masks(*args: ArrayLike, compress: bool = ...) -> list[ArrayLike]: ...
+
+@dataclasses.dataclass
+class BoxplotArtists(collections.abc.Mapping):
+    boxes: Union[list[Line2D], list[Patch]]
+    medians: list[Line2D]
+    whiskers: list[Line2D]
+    caps: list[Line2D]
+    fliers: list[Line2D]
+    means: list[Line2D]
+    def __getitem__(self, key: str) -> Union[list[Line2D], list[Patch]]: ...
+    def __iter__(self) -> Iterator[str]: ...
+    def __len__(self) -> int: ...
+
 def boxplot_stats(
     X: ArrayLike,
     whis: float | tuple[float, float] = ...,
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3645,6 +3645,15 @@ def test_boxplot_mod_artist_after_plotting():
             obj.set_color('green')
 
 
+def test_boxplot_return_dict_and_dataclass():
+    # check that the returned BoxplotArtists works as a dataclass and as a dict
+    # i.e. bp['key'] and bp.key are equivalent
+    fig, ax = plt.subplots()
+    bp = ax.boxplot(np.linspace(0, 1, 11), sym="o")
+    for key in bp:
+        assert bp[key] is getattr(bp, key)
+
+
 @image_comparison(['violinplot_vert_baseline.png',
                    'violinplot_vert_baseline.png'])
 def test_vert_violinplot_baseline():

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+``Boxplot()`` returns a dataclass instead of a dict
	`2`	`+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~`
	`3`	`+The dataclass still supports dict-like access for backwards compatibility.`