test_legend.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644
  1. import collections
  2. import platform
  3. from unittest import mock
  4. import numpy as np
  5. import pytest
  6. from matplotlib.testing.decorators import image_comparison
  7. import matplotlib.pyplot as plt
  8. import matplotlib as mpl
  9. import matplotlib.transforms as mtransforms
  10. import matplotlib.collections as mcollections
  11. from matplotlib.legend_handler import HandlerTuple
  12. import matplotlib.legend as mlegend
  13. from matplotlib import rc_context
  14. def test_legend_ordereddict():
  15. # smoketest that ordereddict inputs work...
  16. X = np.random.randn(10)
  17. Y = np.random.randn(10)
  18. labels = ['a'] * 5 + ['b'] * 5
  19. colors = ['r'] * 5 + ['g'] * 5
  20. fig, ax = plt.subplots()
  21. for x, y, label, color in zip(X, Y, labels, colors):
  22. ax.scatter(x, y, label=label, c=color)
  23. handles, labels = ax.get_legend_handles_labels()
  24. legend = collections.OrderedDict(zip(labels, handles))
  25. ax.legend(legend.values(), legend.keys(),
  26. loc='center left', bbox_to_anchor=(1, .5))
  27. @image_comparison(['legend_auto1'], remove_text=True)
  28. def test_legend_auto1():
  29. """Test automatic legend placement"""
  30. fig = plt.figure()
  31. ax = fig.add_subplot(111)
  32. x = np.arange(100)
  33. ax.plot(x, 50 - x, 'o', label='y=1')
  34. ax.plot(x, x - 50, 'o', label='y=-1')
  35. ax.legend(loc='best')
  36. @image_comparison(['legend_auto2'], remove_text=True)
  37. def test_legend_auto2():
  38. """Test automatic legend placement"""
  39. fig = plt.figure()
  40. ax = fig.add_subplot(111)
  41. x = np.arange(100)
  42. b1 = ax.bar(x, x, align='edge', color='m')
  43. b2 = ax.bar(x, x[::-1], align='edge', color='g')
  44. ax.legend([b1[0], b2[0]], ['up', 'down'], loc='best')
  45. @image_comparison(['legend_auto3'])
  46. def test_legend_auto3():
  47. """Test automatic legend placement"""
  48. fig = plt.figure()
  49. ax = fig.add_subplot(111)
  50. x = [0.9, 0.1, 0.1, 0.9, 0.9, 0.5]
  51. y = [0.95, 0.95, 0.05, 0.05, 0.5, 0.5]
  52. ax.plot(x, y, 'o-', label='line')
  53. ax.set_xlim(0.0, 1.0)
  54. ax.set_ylim(0.0, 1.0)
  55. ax.legend(loc='best')
  56. @image_comparison(['legend_various_labels'], remove_text=True)
  57. def test_various_labels():
  58. # tests all sorts of label types
  59. fig = plt.figure()
  60. ax = fig.add_subplot(121)
  61. ax.plot(np.arange(4), 'o', label=1)
  62. ax.plot(np.linspace(4, 4.1), 'o', label='Développés')
  63. ax.plot(np.arange(4, 1, -1), 'o', label='__nolegend__')
  64. ax.legend(numpoints=1, loc='best')
  65. @image_comparison(['legend_labels_first.png'], remove_text=True)
  66. def test_labels_first():
  67. # test labels to left of markers
  68. fig = plt.figure()
  69. ax = fig.add_subplot(111)
  70. ax.plot(np.arange(10), '-o', label=1)
  71. ax.plot(np.ones(10)*5, ':x', label="x")
  72. ax.plot(np.arange(20, 10, -1), 'd', label="diamond")
  73. ax.legend(loc='best', markerfirst=False)
  74. @image_comparison(['legend_multiple_keys.png'], remove_text=True)
  75. def test_multiple_keys():
  76. # test legend entries with multiple keys
  77. fig = plt.figure()
  78. ax = fig.add_subplot(111)
  79. p1, = ax.plot([1, 2, 3], '-o')
  80. p2, = ax.plot([2, 3, 4], '-x')
  81. p3, = ax.plot([3, 4, 5], '-d')
  82. ax.legend([(p1, p2), (p2, p1), p3], ['two keys', 'pad=0', 'one key'],
  83. numpoints=1,
  84. handler_map={(p1, p2): HandlerTuple(ndivide=None),
  85. (p2, p1): HandlerTuple(ndivide=None, pad=0)})
  86. @image_comparison(['rgba_alpha.png'], remove_text=True,
  87. tol=0 if platform.machine() == 'x86_64' else 0.01)
  88. def test_alpha_rgba():
  89. fig, ax = plt.subplots(1, 1)
  90. ax.plot(range(10), lw=5)
  91. leg = plt.legend(['Longlabel that will go away'], loc='center')
  92. leg.legendPatch.set_facecolor([1, 0, 0, 0.5])
  93. @image_comparison(['rcparam_alpha.png'], remove_text=True,
  94. tol=0 if platform.machine() == 'x86_64' else 0.01)
  95. def test_alpha_rcparam():
  96. fig, ax = plt.subplots(1, 1)
  97. ax.plot(range(10), lw=5)
  98. with mpl.rc_context(rc={'legend.framealpha': .75}):
  99. leg = plt.legend(['Longlabel that will go away'], loc='center')
  100. # this alpha is going to be over-ridden by the rcparam with
  101. # sets the alpha of the patch to be non-None which causes the alpha
  102. # value of the face color to be discarded. This behavior may not be
  103. # ideal, but it is what it is and we should keep track of it changing
  104. leg.legendPatch.set_facecolor([1, 0, 0, 0.5])
  105. @image_comparison(['fancy'], remove_text=True)
  106. def test_fancy():
  107. # using subplot triggers some offsetbox functionality untested elsewhere
  108. plt.subplot(121)
  109. plt.scatter(np.arange(10), np.arange(10, 0, -1), label='XX\nXX')
  110. plt.plot([5] * 10, 'o--', label='XX')
  111. plt.errorbar(np.arange(10), np.arange(10), xerr=0.5,
  112. yerr=0.5, label='XX')
  113. plt.legend(loc="center left", bbox_to_anchor=[1.0, 0.5],
  114. ncol=2, shadow=True, title="My legend", numpoints=1)
  115. @image_comparison(['framealpha'], remove_text=True,
  116. tol=0 if platform.machine() == 'x86_64' else 0.02)
  117. def test_framealpha():
  118. x = np.linspace(1, 100, 100)
  119. y = x
  120. plt.plot(x, y, label='mylabel', lw=10)
  121. plt.legend(framealpha=0.5)
  122. @image_comparison(['scatter_rc3', 'scatter_rc1'], remove_text=True)
  123. def test_rc():
  124. # using subplot triggers some offsetbox functionality untested elsewhere
  125. plt.figure()
  126. ax = plt.subplot(121)
  127. ax.scatter(np.arange(10), np.arange(10, 0, -1), label='three')
  128. ax.legend(loc="center left", bbox_to_anchor=[1.0, 0.5],
  129. title="My legend")
  130. mpl.rcParams['legend.scatterpoints'] = 1
  131. plt.figure()
  132. ax = plt.subplot(121)
  133. ax.scatter(np.arange(10), np.arange(10, 0, -1), label='one')
  134. ax.legend(loc="center left", bbox_to_anchor=[1.0, 0.5],
  135. title="My legend")
  136. @image_comparison(['legend_expand'], remove_text=True)
  137. def test_legend_expand():
  138. """Test expand mode"""
  139. legend_modes = [None, "expand"]
  140. fig, axes_list = plt.subplots(len(legend_modes), 1)
  141. x = np.arange(100)
  142. for ax, mode in zip(axes_list, legend_modes):
  143. ax.plot(x, 50 - x, 'o', label='y=1')
  144. l1 = ax.legend(loc='upper left', mode=mode)
  145. ax.add_artist(l1)
  146. ax.plot(x, x - 50, 'o', label='y=-1')
  147. l2 = ax.legend(loc='right', mode=mode)
  148. ax.add_artist(l2)
  149. ax.legend(loc='lower left', mode=mode, ncol=2)
  150. @image_comparison(['hatching'], remove_text=True, style='default')
  151. def test_hatching():
  152. # Remove this line when this test image is regenerated.
  153. plt.rcParams['text.kerning_factor'] = 6
  154. fig, ax = plt.subplots()
  155. # Patches
  156. patch = plt.Rectangle((0, 0), 0.3, 0.3, hatch='xx',
  157. label='Patch\ndefault color\nfilled')
  158. ax.add_patch(patch)
  159. patch = plt.Rectangle((0.33, 0), 0.3, 0.3, hatch='||', edgecolor='C1',
  160. label='Patch\nexplicit color\nfilled')
  161. ax.add_patch(patch)
  162. patch = plt.Rectangle((0, 0.4), 0.3, 0.3, hatch='xx', fill=False,
  163. label='Patch\ndefault color\nunfilled')
  164. ax.add_patch(patch)
  165. patch = plt.Rectangle((0.33, 0.4), 0.3, 0.3, hatch='||', fill=False,
  166. edgecolor='C1',
  167. label='Patch\nexplicit color\nunfilled')
  168. ax.add_patch(patch)
  169. # Paths
  170. ax.fill_between([0, .15, .3], [.8, .8, .8], [.9, 1.0, .9],
  171. hatch='+', label='Path\ndefault color')
  172. ax.fill_between([.33, .48, .63], [.8, .8, .8], [.9, 1.0, .9],
  173. hatch='+', edgecolor='C2', label='Path\nexplicit color')
  174. ax.set_xlim(-0.01, 1.1)
  175. ax.set_ylim(-0.01, 1.1)
  176. ax.legend(handlelength=4, handleheight=4)
  177. def test_legend_remove():
  178. fig = plt.figure()
  179. ax = fig.add_subplot(1, 1, 1)
  180. lines = ax.plot(range(10))
  181. leg = fig.legend(lines, "test")
  182. leg.remove()
  183. assert fig.legends == []
  184. leg = ax.legend("test")
  185. leg.remove()
  186. assert ax.get_legend() is None
  187. class TestLegendFunction:
  188. # Tests the legend function on the Axes and pyplot.
  189. def test_legend_handle_label(self):
  190. lines = plt.plot(range(10))
  191. with mock.patch('matplotlib.legend.Legend') as Legend:
  192. plt.legend(lines, ['hello world'])
  193. Legend.assert_called_with(plt.gca(), lines, ['hello world'])
  194. def test_legend_handles_only(self):
  195. lines = plt.plot(range(10))
  196. with pytest.raises(TypeError, match='but found an Artist'):
  197. # a single arg is interpreted as labels
  198. # it's a common error to just pass handles
  199. plt.legend(lines)
  200. def test_legend_no_args(self):
  201. lines = plt.plot(range(10), label='hello world')
  202. with mock.patch('matplotlib.legend.Legend') as Legend:
  203. plt.legend()
  204. Legend.assert_called_with(plt.gca(), lines, ['hello world'])
  205. def test_legend_label_args(self):
  206. lines = plt.plot(range(10), label='hello world')
  207. with mock.patch('matplotlib.legend.Legend') as Legend:
  208. plt.legend(['foobar'])
  209. Legend.assert_called_with(plt.gca(), lines, ['foobar'])
  210. def test_legend_three_args(self):
  211. lines = plt.plot(range(10), label='hello world')
  212. with mock.patch('matplotlib.legend.Legend') as Legend:
  213. plt.legend(lines, ['foobar'], loc='right')
  214. Legend.assert_called_with(plt.gca(), lines, ['foobar'], loc='right')
  215. def test_legend_handler_map(self):
  216. lines = plt.plot(range(10), label='hello world')
  217. with mock.patch('matplotlib.legend.'
  218. '_get_legend_handles_labels') as handles_labels:
  219. handles_labels.return_value = lines, ['hello world']
  220. plt.legend(handler_map={'1': 2})
  221. handles_labels.assert_called_with([plt.gca()], {'1': 2})
  222. def test_kwargs(self):
  223. fig, ax = plt.subplots(1, 1)
  224. th = np.linspace(0, 2*np.pi, 1024)
  225. lns, = ax.plot(th, np.sin(th), label='sin', lw=5)
  226. lnc, = ax.plot(th, np.cos(th), label='cos', lw=5)
  227. with mock.patch('matplotlib.legend.Legend') as Legend:
  228. ax.legend(labels=('a', 'b'), handles=(lnc, lns))
  229. Legend.assert_called_with(ax, (lnc, lns), ('a', 'b'))
  230. def test_warn_args_kwargs(self):
  231. fig, ax = plt.subplots(1, 1)
  232. th = np.linspace(0, 2*np.pi, 1024)
  233. lns, = ax.plot(th, np.sin(th), label='sin', lw=5)
  234. lnc, = ax.plot(th, np.cos(th), label='cos', lw=5)
  235. with pytest.warns(UserWarning) as record:
  236. ax.legend((lnc, lns), labels=('a', 'b'))
  237. assert len(record) == 1
  238. assert str(record[0].message) == (
  239. "You have mixed positional and keyword arguments, some input may "
  240. "be discarded.")
  241. def test_parasite(self):
  242. from mpl_toolkits.axes_grid1 import host_subplot
  243. host = host_subplot(111)
  244. par = host.twinx()
  245. p1, = host.plot([0, 1, 2], [0, 1, 2], label="Density")
  246. p2, = par.plot([0, 1, 2], [0, 3, 2], label="Temperature")
  247. with mock.patch('matplotlib.legend.Legend') as Legend:
  248. plt.legend()
  249. Legend.assert_called_with(host, [p1, p2], ['Density', 'Temperature'])
  250. class TestLegendFigureFunction:
  251. # Tests the legend function for figure
  252. def test_legend_handle_label(self):
  253. fig, ax = plt.subplots()
  254. lines = ax.plot(range(10))
  255. with mock.patch('matplotlib.legend.Legend') as Legend:
  256. fig.legend(lines, ['hello world'])
  257. Legend.assert_called_with(fig, lines, ['hello world'],
  258. bbox_transform=fig.transFigure)
  259. def test_legend_no_args(self):
  260. fig, ax = plt.subplots()
  261. lines = ax.plot(range(10), label='hello world')
  262. with mock.patch('matplotlib.legend.Legend') as Legend:
  263. fig.legend()
  264. Legend.assert_called_with(fig, lines, ['hello world'],
  265. bbox_transform=fig.transFigure)
  266. def test_legend_label_arg(self):
  267. fig, ax = plt.subplots()
  268. lines = ax.plot(range(10))
  269. with mock.patch('matplotlib.legend.Legend') as Legend:
  270. fig.legend(['foobar'])
  271. Legend.assert_called_with(fig, lines, ['foobar'],
  272. bbox_transform=fig.transFigure)
  273. def test_legend_label_three_args(self):
  274. fig, ax = plt.subplots()
  275. lines = ax.plot(range(10))
  276. with mock.patch('matplotlib.legend.Legend') as Legend:
  277. fig.legend(lines, ['foobar'], 'right')
  278. Legend.assert_called_with(fig, lines, ['foobar'], 'right',
  279. bbox_transform=fig.transFigure)
  280. def test_legend_label_three_args_pluskw(self):
  281. # test that third argument and loc= called together give
  282. # Exception
  283. fig, ax = plt.subplots()
  284. lines = ax.plot(range(10))
  285. with pytest.raises(Exception):
  286. fig.legend(lines, ['foobar'], 'right', loc='left')
  287. def test_legend_kw_args(self):
  288. fig, axs = plt.subplots(1, 2)
  289. lines = axs[0].plot(range(10))
  290. lines2 = axs[1].plot(np.arange(10) * 2.)
  291. with mock.patch('matplotlib.legend.Legend') as Legend:
  292. fig.legend(loc='right', labels=('a', 'b'), handles=(lines, lines2))
  293. Legend.assert_called_with(
  294. fig, (lines, lines2), ('a', 'b'), loc='right',
  295. bbox_transform=fig.transFigure)
  296. def test_warn_args_kwargs(self):
  297. fig, axs = plt.subplots(1, 2)
  298. lines = axs[0].plot(range(10))
  299. lines2 = axs[1].plot(np.arange(10) * 2.)
  300. with pytest.warns(UserWarning) as record:
  301. fig.legend((lines, lines2), labels=('a', 'b'))
  302. assert len(record) == 1
  303. assert str(record[0].message) == (
  304. "You have mixed positional and keyword arguments, some input may "
  305. "be discarded.")
  306. @image_comparison(['legend_stackplot.png'])
  307. def test_legend_stackplot():
  308. """Test legend for PolyCollection using stackplot."""
  309. # related to #1341, #1943, and PR #3303
  310. fig = plt.figure()
  311. ax = fig.add_subplot(111)
  312. x = np.linspace(0, 10, 10)
  313. y1 = 1.0 * x
  314. y2 = 2.0 * x + 1
  315. y3 = 3.0 * x + 2
  316. ax.stackplot(x, y1, y2, y3, labels=['y1', 'y2', 'y3'])
  317. ax.set_xlim((0, 10))
  318. ax.set_ylim((0, 70))
  319. ax.legend(loc='best')
  320. def test_cross_figure_patch_legend():
  321. fig, ax = plt.subplots()
  322. fig2, ax2 = plt.subplots()
  323. brs = ax.bar(range(3), range(3))
  324. fig2.legend(brs, 'foo')
  325. def test_nanscatter():
  326. fig, ax = plt.subplots()
  327. h = ax.scatter([np.nan], [np.nan], marker="o",
  328. facecolor="r", edgecolor="r", s=3)
  329. ax.legend([h], ["scatter"])
  330. fig, ax = plt.subplots()
  331. for color in ['red', 'green', 'blue']:
  332. n = 750
  333. x, y = np.random.rand(2, n)
  334. scale = 200.0 * np.random.rand(n)
  335. ax.scatter(x, y, c=color, s=scale, label=color,
  336. alpha=0.3, edgecolors='none')
  337. ax.legend()
  338. ax.grid(True)
  339. def test_legend_repeatcheckok():
  340. fig, ax = plt.subplots()
  341. ax.scatter(0.0, 1.0, color='k', marker='o', label='test')
  342. ax.scatter(0.5, 0.0, color='r', marker='v', label='test')
  343. ax.legend()
  344. hand, lab = mlegend._get_legend_handles_labels([ax])
  345. assert len(lab) == 2
  346. fig, ax = plt.subplots()
  347. ax.scatter(0.0, 1.0, color='k', marker='o', label='test')
  348. ax.scatter(0.5, 0.0, color='k', marker='v', label='test')
  349. ax.legend()
  350. hand, lab = mlegend._get_legend_handles_labels([ax])
  351. assert len(lab) == 2
  352. @image_comparison(['not_covering_scatter.png'])
  353. def test_not_covering_scatter():
  354. colors = ['b', 'g', 'r']
  355. for n in range(3):
  356. plt.scatter([n], [n], color=colors[n])
  357. plt.legend(['foo', 'foo', 'foo'], loc='best')
  358. plt.gca().set_xlim(-0.5, 2.2)
  359. plt.gca().set_ylim(-0.5, 2.2)
  360. @image_comparison(['not_covering_scatter_transform.png'])
  361. def test_not_covering_scatter_transform():
  362. # Offsets point to top left, the default auto position
  363. offset = mtransforms.Affine2D().translate(-20, 20)
  364. x = np.linspace(0, 30, 1000)
  365. plt.plot(x, x)
  366. plt.scatter([20], [10], transform=offset + plt.gca().transData)
  367. plt.legend(['foo', 'bar'], loc='best')
  368. def test_linecollection_scaled_dashes():
  369. lines1 = [[(0, .5), (.5, 1)], [(.3, .6), (.2, .2)]]
  370. lines2 = [[[0.7, .2], [.8, .4]], [[.5, .7], [.6, .1]]]
  371. lines3 = [[[0.6, .2], [.8, .4]], [[.5, .7], [.1, .1]]]
  372. lc1 = mcollections.LineCollection(lines1, linestyles="--", lw=3)
  373. lc2 = mcollections.LineCollection(lines2, linestyles="-.")
  374. lc3 = mcollections.LineCollection(lines3, linestyles=":", lw=.5)
  375. fig, ax = plt.subplots()
  376. ax.add_collection(lc1)
  377. ax.add_collection(lc2)
  378. ax.add_collection(lc3)
  379. leg = ax.legend([lc1, lc2, lc3], ["line1", "line2", 'line 3'])
  380. h1, h2, h3 = leg.legendHandles
  381. for oh, lh in zip((lc1, lc2, lc3), (h1, h2, h3)):
  382. assert oh.get_linestyles()[0][1] == lh._dashSeq
  383. assert oh.get_linestyles()[0][0] == lh._dashOffset
  384. def test_handler_numpoints():
  385. """Test legend handler with numpoints <= 1."""
  386. # related to #6921 and PR #8478
  387. fig, ax = plt.subplots()
  388. ax.plot(range(5), label='test')
  389. ax.legend(numpoints=0.5)
  390. def test_empty_bar_chart_with_legend():
  391. """Test legend when bar chart is empty with a label."""
  392. # related to issue #13003. Calling plt.legend() should not
  393. # raise an IndexError.
  394. plt.bar([], [], label='test')
  395. plt.legend()
  396. def test_shadow_framealpha():
  397. # Test if framealpha is activated when shadow is True
  398. # and framealpha is not explicitly passed'''
  399. fig, ax = plt.subplots()
  400. ax.plot(range(100), label="test")
  401. leg = ax.legend(shadow=True, facecolor='w')
  402. assert leg.get_frame().get_alpha() == 1
  403. def test_legend_title_empty():
  404. # test that if we don't set the legend title, that
  405. # it comes back as an empty string, and that it is not
  406. # visible:
  407. fig, ax = plt.subplots()
  408. ax.plot(range(10))
  409. leg = ax.legend()
  410. assert leg.get_title().get_text() == ""
  411. assert not leg.get_title().get_visible()
  412. def test_legend_proper_window_extent():
  413. # test that legend returns the expected extent under various dpi...
  414. fig, ax = plt.subplots(dpi=100)
  415. ax.plot(range(10), label='Aardvark')
  416. leg = ax.legend()
  417. x01 = leg.get_window_extent(fig.canvas.get_renderer()).x0
  418. fig, ax = plt.subplots(dpi=200)
  419. ax.plot(range(10), label='Aardvark')
  420. leg = ax.legend()
  421. x02 = leg.get_window_extent(fig.canvas.get_renderer()).x0
  422. assert pytest.approx(x01*2, 0.1) == x02
  423. def test_window_extent_cached_renderer():
  424. fig, ax = plt.subplots(dpi=100)
  425. ax.plot(range(10), label='Aardvark')
  426. leg = ax.legend()
  427. leg2 = fig.legend()
  428. fig.canvas.draw()
  429. # check that get_window_extent will use the cached renderer
  430. leg.get_window_extent()
  431. leg2.get_window_extent()
  432. def test_legend_title_fontsize():
  433. # test the title_fontsize kwarg
  434. fig, ax = plt.subplots()
  435. ax.plot(range(10))
  436. leg = ax.legend(title='Aardvark', title_fontsize=22)
  437. assert leg.get_title().get_fontsize() == 22
  438. def test_legend_labelcolor_single():
  439. # test labelcolor for a single color
  440. fig, ax = plt.subplots()
  441. ax.plot(np.arange(10), np.arange(10)*1, label='#1')
  442. ax.plot(np.arange(10), np.arange(10)*2, label='#2')
  443. ax.plot(np.arange(10), np.arange(10)*3, label='#3')
  444. leg = ax.legend(labelcolor='red')
  445. for text in leg.get_texts():
  446. assert mpl.colors.same_color(text.get_color(), 'red')
  447. def test_legend_labelcolor_list():
  448. # test labelcolor for a list of colors
  449. fig, ax = plt.subplots()
  450. ax.plot(np.arange(10), np.arange(10)*1, label='#1')
  451. ax.plot(np.arange(10), np.arange(10)*2, label='#2')
  452. ax.plot(np.arange(10), np.arange(10)*3, label='#3')
  453. leg = ax.legend(labelcolor=['r', 'g', 'b'])
  454. for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
  455. assert mpl.colors.same_color(text.get_color(), color)
  456. def test_legend_labelcolor_linecolor():
  457. # test the labelcolor for labelcolor='linecolor'
  458. fig, ax = plt.subplots()
  459. ax.plot(np.arange(10), np.arange(10)*1, label='#1', color='r')
  460. ax.plot(np.arange(10), np.arange(10)*2, label='#2', color='g')
  461. ax.plot(np.arange(10), np.arange(10)*3, label='#3', color='b')
  462. leg = ax.legend(labelcolor='linecolor')
  463. for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
  464. assert mpl.colors.same_color(text.get_color(), color)
  465. def test_legend_labelcolor_markeredgecolor():
  466. # test the labelcolor for labelcolor='markeredgecolor'
  467. fig, ax = plt.subplots()
  468. ax.plot(np.arange(10), np.arange(10)*1, label='#1', markeredgecolor='r')
  469. ax.plot(np.arange(10), np.arange(10)*2, label='#2', markeredgecolor='g')
  470. ax.plot(np.arange(10), np.arange(10)*3, label='#3', markeredgecolor='b')
  471. leg = ax.legend(labelcolor='markeredgecolor')
  472. for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
  473. assert mpl.colors.same_color(text.get_color(), color)
  474. def test_legend_labelcolor_markerfacecolor():
  475. # test the labelcolor for labelcolor='markerfacecolor'
  476. fig, ax = plt.subplots()
  477. ax.plot(np.arange(10), np.arange(10)*1, label='#1', markerfacecolor='r')
  478. ax.plot(np.arange(10), np.arange(10)*2, label='#2', markerfacecolor='g')
  479. ax.plot(np.arange(10), np.arange(10)*3, label='#3', markerfacecolor='b')
  480. leg = ax.legend(labelcolor='markerfacecolor')
  481. for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
  482. assert mpl.colors.same_color(text.get_color(), color)
  483. def test_get_set_draggable():
  484. legend = plt.legend()
  485. assert not legend.get_draggable()
  486. legend.set_draggable(True)
  487. assert legend.get_draggable()
  488. legend.set_draggable(False)
  489. assert not legend.get_draggable()
  490. def test_alpha_handles():
  491. x, n, hh = plt.hist([1, 2, 3], alpha=0.25, label='data', color='red')
  492. legend = plt.legend()
  493. for lh in legend.legendHandles:
  494. lh.set_alpha(1.0)
  495. assert lh.get_facecolor()[:-1] == hh[1].get_facecolor()[:-1]
  496. assert lh.get_edgecolor()[:-1] == hh[1].get_edgecolor()[:-1]
  497. def test_warn_big_data_best_loc():
  498. fig, ax = plt.subplots()
  499. fig.canvas.draw() # So that we can call draw_artist later.
  500. for idx in range(1000):
  501. ax.plot(np.arange(5000), label=idx)
  502. with rc_context({'legend.loc': 'best'}):
  503. legend = ax.legend()
  504. with pytest.warns(UserWarning) as records:
  505. fig.draw_artist(legend) # Don't bother drawing the lines -- it's slow.
  506. # The _find_best_position method of Legend is called twice, duplicating
  507. # the warning message.
  508. assert len(records) == 2
  509. for record in records:
  510. assert str(record.message) == (
  511. 'Creating legend with loc="best" can be slow with large '
  512. 'amounts of data.')
  513. def test_no_warn_big_data_when_loc_specified():
  514. fig, ax = plt.subplots()
  515. fig.canvas.draw()
  516. for idx in range(1000):
  517. ax.plot(np.arange(5000), label=idx)
  518. legend = ax.legend('best')
  519. fig.draw_artist(legend) # Check that no warning is emitted.