Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 2c99922

Browse files
committed
legend supports Bar and Errorbar plots
1 parent fe03e22 commit 2c99922

5 files changed

Lines changed: 247 additions & 63 deletions

File tree

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import matplotlib.pyplot as plt
2+
3+
if 1:
4+
5+
ax = plt.subplot(111)
6+
7+
b1 = ax.bar([0, 1, 2], [0.2, 0.3, 0.1], width=0.4,
8+
label="Bar 1", align="center")
9+
10+
b2 = ax.bar([0.5, 1.5, 2.5], [0.3, 0.2, 0.2], color="red", width=0.4,
11+
label="Bar 2", align="center")
12+
13+
err1 = ax.errorbar([0, 1, 2], [2, 3, 1], xerr=0.4, fmt="s",
14+
label="test 1")
15+
err2 = ax.errorbar([0, 1, 2], [3, 2, 4], yerr=0.3, fmt="o",
16+
label="test 2")
17+
err3 = ax.errorbar([0, 1, 2], [1, 1, 3], xerr=0.4, yerr=0.3, fmt="^",
18+
label="test 3")
19+
20+
# legend
21+
leg1 = plt.legend(loc=1)
22+
23+
# legend of selected artists
24+
artists = [b1, err2]
25+
leg2 = plt.legend(artists, [a.get_label() for a in artists], loc=2)
26+
27+
# custome handler
28+
import matplotlib.legend_handler as mlegend_handler
29+
myhandler = mlegend_handler.HandlerErrorbar(npoints=1)
30+
31+
leg3 = plt.legend([err1, err3], ["T1", "T2"], loc=3,
32+
handler_map={err3:myhandler})
33+
34+
plt.gca().add_artist(leg1)
35+
plt.gca().add_artist(leg2)
36+
37+
plt.show()
38+

lib/matplotlib/axes.py

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
import matplotlib.transforms as mtransforms
3535
import matplotlib.tri as mtri
3636

37-
from matplotlib.container import BarContainer
37+
from matplotlib.container import BarContainer, ErrorbarContainer
3838

3939
iterable = cbook.iterable
4040
is_string_like = cbook.is_string_like
@@ -1507,6 +1507,7 @@ def add_container(self, container):
15071507
if not label:
15081508
container.set_label('_container%d'%len(self.containers))
15091509
self.containers.append(container)
1510+
container.set_remove_method(lambda h: self.containers.remove(container))
15101511
return container
15111512

15121513

@@ -4200,19 +4201,25 @@ def xcorr(self, x, y, normed=True, detrend=mlab.detrend_none,
42004201

42014202
def _get_legend_handles(self):
42024203
"return artists that will be used as handles for legend"
4203-
handles = self.lines[:]
4204-
handles.extend(self.patches)
4204+
handles_original = self.lines + self.patches + \
4205+
self.collections + self.containers
42054206

42064207
# collections
42074208
legend_map_keys = mlegend.Legend._default_handler_map.keys()
4208-
collection_types = [cls for cls in legend_map_keys \
4209-
if issubclass(cls, mcoll.Collection)]
4209+
#collection_types = [cls for cls in legend_map_keys \
4210+
# if issubclass(cls, mcoll.Collection)]
42104211

4211-
for cls in collection_types:
4212-
handles.extend([c for c in self.collections
4213-
if isinstance(c, cls)])
4212+
handles = []
4213+
for h in handles_original:
4214+
if h.get_label().startswith('_'):
4215+
continue
4216+
4217+
# check subclass
4218+
for cls in legend_map_keys:
4219+
if isinstance(h, cls):
4220+
handles.append(h)
4221+
break
42144222

4215-
handles.extend([c for c in self.containers])
42164223
return handles
42174224

42184225

@@ -4231,8 +4238,7 @@ def get_legend_handles_labels(self):
42314238
labels = []
42324239
for handle in self._get_legend_handles():
42334240
label = handle.get_label()
4234-
if (label is not None and
4235-
label != '' and not label.startswith('_')):
4241+
if (label is not None and label != ''):
42364242
handles.append(handle)
42374243
labels.append(label)
42384244

@@ -5140,6 +5146,8 @@ def errorbar(self, x, y, yerr=None, xerr=None,
51405146
holdstate = self._hold
51415147
self._hold = True
51425148

5149+
label = kwargs.pop("label", None)
5150+
51435151
# make sure all the args are iterable; use lists not arrays to
51445152
# preserve units
51455153
if not iterable(x):
@@ -5159,7 +5167,7 @@ def errorbar(self, x, y, yerr=None, xerr=None,
51595167
l0 = None
51605168

51615169
if barsabove and fmt is not None:
5162-
l0, = self.plot(x,y,fmt,**kwargs)
5170+
l0, = self.plot(x,y,fmt,label="_nolegend_", **kwargs)
51635171

51645172
barcols = []
51655173
caplines = []
@@ -5312,7 +5320,14 @@ def xywhere(xs, ys, mask):
53125320

53135321
self.autoscale_view()
53145322
self._hold = holdstate
5315-
return (l0, caplines, barcols)
5323+
5324+
errorbar_container = ErrorbarContainer((l0, tuple(caplines), tuple(barcols)),
5325+
has_xerr=(xerr is not None),
5326+
has_yerr=(yerr is not None),
5327+
label=label)
5328+
self.containers.append(errorbar_container)
5329+
5330+
return errorbar_container # (l0, caplines, barcols)
53165331

53175332
def boxplot(self, x, notch=0, sym='b+', vert=1, whis=1.5,
53185333
positions=None, widths=None, patch_artist=False,

lib/matplotlib/container.py

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import matplotlib.cbook as cbook
12

23
class Container(tuple):
34
"""
@@ -16,8 +17,19 @@ def __init__(self, kl, label=None):
1617
self._oid = 0 # an observer id
1718
self._propobservers = {} # a dict from oids to funcs
1819

20+
self._remove_method = None
21+
1922
self.set_label(label)
2023

24+
def set_remove_method(self, f):
25+
self._remove_method = f
26+
27+
def remove(self):
28+
for c in self:
29+
c.remove()
30+
31+
if self._remove_method:
32+
self._remove_method()
2133

2234
def get_label(self):
2335
"""
@@ -68,9 +80,8 @@ def pchanged(self):
6880
for oid, func in self._propobservers.items():
6981
func(self)
7082

71-
def remove(self):
72-
for c in self:
73-
c.remove()
83+
def get_children(self):
84+
return list(cbook.flatten(self))
7485

7586

7687
class BarContainer(Container):
@@ -81,14 +92,12 @@ def __init__(self, patches, errorbar=None, **kwargs):
8192
Container.__init__(self, patches, **kwargs)
8293

8394

84-
if __name__ == '__main__':
85-
import matplotlib.pyplot as plt
86-
87-
plt.clf()
88-
bb1 = plt.bar([0, 1, 2], [2, 3, 1], label="test", width=0.4)
89-
bb2 = plt.bar([0.5, 1.5, 2.5], [2, 3, 1], label="test2", color="red", width=0.4)
90-
#cont1 = Container(bb1, err=3)
95+
class ErrorbarContainer(Container):
9196

97+
def __init__(self, lines, has_xerr=False, has_yerr=False, **kwargs):
98+
self.lines = lines
99+
self.has_xerr = has_xerr
100+
self.has_yerr = has_yerr
101+
Container.__init__(self, lines, **kwargs)
92102

93-
#aa = BarContainer(("a",), errorbar=1)
94103

lib/matplotlib/legend.py

Lines changed: 26 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737

3838
from matplotlib.offsetbox import HPacker, VPacker, TextArea, DrawingArea, DraggableOffsetBox
3939

40+
from matplotlib.container import ErrorbarContainer, BarContainer
4041
import legend_handler
4142

4243

@@ -69,7 +70,7 @@ def finalize_offset(self):
6970
self._update_bbox_to_anchor(loc_in_canvas)
7071
else:
7172
raise RuntimeError("update parameter '%s' is not supported." % self.update)
72-
73+
7374
def _update_loc(self, loc_in_canvas):
7475
bbox = self.legend.get_bbox_to_anchor()
7576

@@ -78,7 +79,7 @@ def _update_loc(self, loc_in_canvas):
7879
if bbox.width ==0 or bbox.height ==0:
7980
self.legend.set_bbox_to_anchor(None)
8081
bbox = self.legend.get_bbox_to_anchor()
81-
82+
8283
_bbox_transform = BboxTransformFrom(bbox)
8384
self.legend._loc = tuple(_bbox_transform.transform_point(loc_in_canvas))
8485

@@ -476,12 +477,15 @@ def _approx_text_height(self, renderer=None):
476477
return renderer.points_to_pixels(self._fontsize)
477478

478479

479-
_default_handler_map = {Line2D:legend_handler.HandlerLine2D(),
480-
Patch:legend_handler.HandlerPatch(),
481-
LineCollection:legend_handler.HandlerLineCollection(),
482-
RegularPolyCollection:legend_handler.HandlerRegularPolyCollection(),
483-
CircleCollection:legend_handler.HandlerCircleCollection()
484-
}
480+
_default_handler_map = {
481+
ErrorbarContainer:legend_handler.HandlerErrorbar(),
482+
Line2D:legend_handler.HandlerLine2D(),
483+
Patch:legend_handler.HandlerPatch(),
484+
LineCollection:legend_handler.HandlerLineCollection(),
485+
RegularPolyCollection:legend_handler.HandlerRegularPolyCollection(),
486+
CircleCollection:legend_handler.HandlerCircleCollection(),
487+
BarContainer:legend_handler.HandlerPatch(update_func=legend_handler.update_from_first_child)
488+
}
485489

486490
def get_legend_handler_map(self):
487491
if self._handler_map:
@@ -492,20 +496,21 @@ def get_legend_handler_map(self):
492496
return self._default_handler_map
493497

494498
def get_legend_handler(self, legend_handler_map, orig_handle):
495-
if orig_handle in legend_handler_map:
496-
handler = legend_handler_map[orig_handle]
499+
legend_handler_keys = legend_handler_map.keys()
500+
if orig_handle in legend_handler_keys:
501+
handler = legend_handler_map[orig_handle]
502+
else:
503+
504+
for handle_type in type(orig_handle).mro():
505+
if handle_type in legend_handler_map:
506+
handler = legend_handler_map[handle_type]
507+
break
497508
else:
498-
499-
for handle_type in type(orig_handle).mro():
500-
if handle_type in legend_handler_map:
501-
handler = legend_handler_map[handle_type]
502-
break
503-
else:
504-
handler = None
509+
handler = None
510+
511+
return handler
505512

506-
return handler
507513

508-
509514
def _init_legend_box(self, handles, labels):
510515
"""
511516
Initiallize the legend_box. The legend_box is an instance of
@@ -562,7 +567,7 @@ def _init_legend_box(self, handles, labels):
562567
warnings.warn("Legend does not support %s\nUse proxy artist instead.\n\nhttp://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist\n" % (str(orig_handle),))
563568
handle_list.append(None)
564569
continue
565-
570+
566571
textbox = TextArea(lab, textprops=label_prop,
567572
multilinebaseline=True, minimumdescent=True)
568573
text_list.append(textbox._text)
@@ -578,7 +583,7 @@ def _init_legend_box(self, handles, labels):
578583
fontsize,
579584
handlebox)
580585
handle_list.append(handle)
581-
586+
582587

583588

584589

0 commit comments

Comments
 (0)