test_matplotlib.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. DataFrame,
  5. IndexSlice,
  6. Series,
  7. )
  8. pytest.importorskip("matplotlib")
  9. pytest.importorskip("jinja2")
  10. import matplotlib as mpl
  11. from pandas.io.formats.style import Styler
  12. @pytest.fixture
  13. def df():
  14. return DataFrame([[1, 2], [2, 4]], columns=["A", "B"])
  15. @pytest.fixture
  16. def styler(df):
  17. return Styler(df, uuid_len=0)
  18. @pytest.fixture
  19. def df_blank():
  20. return DataFrame([[0, 0], [0, 0]], columns=["A", "B"], index=["X", "Y"])
  21. @pytest.fixture
  22. def styler_blank(df_blank):
  23. return Styler(df_blank, uuid_len=0)
  24. @pytest.mark.parametrize("f", ["background_gradient", "text_gradient"])
  25. def test_function_gradient(styler, f):
  26. for c_map in [None, "YlOrRd"]:
  27. result = getattr(styler, f)(cmap=c_map)._compute().ctx
  28. assert all("#" in x[0][1] for x in result.values())
  29. assert result[(0, 0)] == result[(0, 1)]
  30. assert result[(1, 0)] == result[(1, 1)]
  31. @pytest.mark.parametrize("f", ["background_gradient", "text_gradient"])
  32. def test_background_gradient_color(styler, f):
  33. result = getattr(styler, f)(subset=IndexSlice[1, "A"])._compute().ctx
  34. if f == "background_gradient":
  35. assert result[(1, 0)] == [("background-color", "#fff7fb"), ("color", "#000000")]
  36. elif f == "text_gradient":
  37. assert result[(1, 0)] == [("color", "#fff7fb")]
  38. @pytest.mark.parametrize(
  39. "axis, expected",
  40. [
  41. (0, ["low", "low", "high", "high"]),
  42. (1, ["low", "high", "low", "high"]),
  43. (None, ["low", "mid", "mid", "high"]),
  44. ],
  45. )
  46. @pytest.mark.parametrize("f", ["background_gradient", "text_gradient"])
  47. def test_background_gradient_axis(styler, axis, expected, f):
  48. if f == "background_gradient":
  49. colors = {
  50. "low": [("background-color", "#f7fbff"), ("color", "#000000")],
  51. "mid": [("background-color", "#abd0e6"), ("color", "#000000")],
  52. "high": [("background-color", "#08306b"), ("color", "#f1f1f1")],
  53. }
  54. elif f == "text_gradient":
  55. colors = {
  56. "low": [("color", "#f7fbff")],
  57. "mid": [("color", "#abd0e6")],
  58. "high": [("color", "#08306b")],
  59. }
  60. result = getattr(styler, f)(cmap="Blues", axis=axis)._compute().ctx
  61. for i, cell in enumerate([(0, 0), (0, 1), (1, 0), (1, 1)]):
  62. assert result[cell] == colors[expected[i]]
  63. @pytest.mark.parametrize(
  64. "cmap, expected",
  65. [
  66. (
  67. "PuBu",
  68. {
  69. (4, 5): [("background-color", "#86b0d3"), ("color", "#000000")],
  70. (4, 6): [("background-color", "#83afd3"), ("color", "#f1f1f1")],
  71. },
  72. ),
  73. (
  74. "YlOrRd",
  75. {
  76. (4, 8): [("background-color", "#fd913e"), ("color", "#000000")],
  77. (4, 9): [("background-color", "#fd8f3d"), ("color", "#f1f1f1")],
  78. },
  79. ),
  80. (
  81. None,
  82. {
  83. (7, 0): [("background-color", "#48c16e"), ("color", "#f1f1f1")],
  84. (7, 1): [("background-color", "#4cc26c"), ("color", "#000000")],
  85. },
  86. ),
  87. ],
  88. )
  89. def test_text_color_threshold(cmap, expected):
  90. # GH 39888
  91. df = DataFrame(np.arange(100).reshape(10, 10))
  92. result = df.style.background_gradient(cmap=cmap, axis=None)._compute().ctx
  93. for k in expected.keys():
  94. assert result[k] == expected[k]
  95. def test_background_gradient_vmin_vmax():
  96. # GH 12145
  97. df = DataFrame(range(5))
  98. ctx = df.style.background_gradient(vmin=1, vmax=3)._compute().ctx
  99. assert ctx[(0, 0)] == ctx[(1, 0)]
  100. assert ctx[(4, 0)] == ctx[(3, 0)]
  101. def test_background_gradient_int64():
  102. # GH 28869
  103. df1 = Series(range(3)).to_frame()
  104. df2 = Series(range(3), dtype="Int64").to_frame()
  105. ctx1 = df1.style.background_gradient()._compute().ctx
  106. ctx2 = df2.style.background_gradient()._compute().ctx
  107. assert ctx2[(0, 0)] == ctx1[(0, 0)]
  108. assert ctx2[(1, 0)] == ctx1[(1, 0)]
  109. assert ctx2[(2, 0)] == ctx1[(2, 0)]
  110. @pytest.mark.parametrize(
  111. "axis, gmap, expected",
  112. [
  113. (
  114. 0,
  115. [1, 2],
  116. {
  117. (0, 0): [("background-color", "#fff7fb"), ("color", "#000000")],
  118. (1, 0): [("background-color", "#023858"), ("color", "#f1f1f1")],
  119. (0, 1): [("background-color", "#fff7fb"), ("color", "#000000")],
  120. (1, 1): [("background-color", "#023858"), ("color", "#f1f1f1")],
  121. },
  122. ),
  123. (
  124. 1,
  125. [1, 2],
  126. {
  127. (0, 0): [("background-color", "#fff7fb"), ("color", "#000000")],
  128. (1, 0): [("background-color", "#fff7fb"), ("color", "#000000")],
  129. (0, 1): [("background-color", "#023858"), ("color", "#f1f1f1")],
  130. (1, 1): [("background-color", "#023858"), ("color", "#f1f1f1")],
  131. },
  132. ),
  133. (
  134. None,
  135. np.array([[2, 1], [1, 2]]),
  136. {
  137. (0, 0): [("background-color", "#023858"), ("color", "#f1f1f1")],
  138. (1, 0): [("background-color", "#fff7fb"), ("color", "#000000")],
  139. (0, 1): [("background-color", "#fff7fb"), ("color", "#000000")],
  140. (1, 1): [("background-color", "#023858"), ("color", "#f1f1f1")],
  141. },
  142. ),
  143. ],
  144. )
  145. def test_background_gradient_gmap_array(styler_blank, axis, gmap, expected):
  146. # tests when gmap is given as a sequence and converted to ndarray
  147. result = styler_blank.background_gradient(axis=axis, gmap=gmap)._compute().ctx
  148. assert result == expected
  149. @pytest.mark.parametrize(
  150. "gmap, axis", [([1, 2, 3], 0), ([1, 2], 1), (np.array([[1, 2], [1, 2]]), None)]
  151. )
  152. def test_background_gradient_gmap_array_raises(gmap, axis):
  153. # test when gmap as converted ndarray is bad shape
  154. df = DataFrame([[0, 0, 0], [0, 0, 0]])
  155. msg = "supplied 'gmap' is not correct shape"
  156. with pytest.raises(ValueError, match=msg):
  157. df.style.background_gradient(gmap=gmap, axis=axis)._compute()
  158. @pytest.mark.parametrize(
  159. "gmap",
  160. [
  161. DataFrame( # reverse the columns
  162. [[2, 1], [1, 2]], columns=["B", "A"], index=["X", "Y"]
  163. ),
  164. DataFrame( # reverse the index
  165. [[2, 1], [1, 2]], columns=["A", "B"], index=["Y", "X"]
  166. ),
  167. DataFrame( # reverse the index and columns
  168. [[1, 2], [2, 1]], columns=["B", "A"], index=["Y", "X"]
  169. ),
  170. DataFrame( # add unnecessary columns
  171. [[1, 2, 3], [2, 1, 3]], columns=["A", "B", "C"], index=["X", "Y"]
  172. ),
  173. DataFrame( # add unnecessary index
  174. [[1, 2], [2, 1], [3, 3]], columns=["A", "B"], index=["X", "Y", "Z"]
  175. ),
  176. ],
  177. )
  178. @pytest.mark.parametrize(
  179. "subset, exp_gmap", # exp_gmap is underlying map DataFrame should conform to
  180. [
  181. (None, [[1, 2], [2, 1]]),
  182. (["A"], [[1], [2]]), # slice only column "A" in data and gmap
  183. (["B", "A"], [[2, 1], [1, 2]]), # reverse the columns in data
  184. (IndexSlice["X", :], [[1, 2]]), # slice only index "X" in data and gmap
  185. (IndexSlice[["Y", "X"], :], [[2, 1], [1, 2]]), # reverse the index in data
  186. ],
  187. )
  188. def test_background_gradient_gmap_dataframe_align(styler_blank, gmap, subset, exp_gmap):
  189. # test gmap given as DataFrame that it aligns to the data including subset
  190. expected = styler_blank.background_gradient(axis=None, gmap=exp_gmap, subset=subset)
  191. result = styler_blank.background_gradient(axis=None, gmap=gmap, subset=subset)
  192. assert expected._compute().ctx == result._compute().ctx
  193. @pytest.mark.parametrize(
  194. "gmap, axis, exp_gmap",
  195. [
  196. (Series([2, 1], index=["Y", "X"]), 0, [[1, 1], [2, 2]]), # revrse the index
  197. (Series([2, 1], index=["B", "A"]), 1, [[1, 2], [1, 2]]), # revrse the cols
  198. (Series([1, 2, 3], index=["X", "Y", "Z"]), 0, [[1, 1], [2, 2]]), # add idx
  199. (Series([1, 2, 3], index=["A", "B", "C"]), 1, [[1, 2], [1, 2]]), # add col
  200. ],
  201. )
  202. def test_background_gradient_gmap_series_align(styler_blank, gmap, axis, exp_gmap):
  203. # test gmap given as Series that it aligns to the data including subset
  204. expected = styler_blank.background_gradient(axis=None, gmap=exp_gmap)._compute()
  205. result = styler_blank.background_gradient(axis=axis, gmap=gmap)._compute()
  206. assert expected.ctx == result.ctx
  207. @pytest.mark.parametrize(
  208. "gmap, axis",
  209. [
  210. (DataFrame([[1, 2], [2, 1]], columns=["A", "B"], index=["X", "Y"]), 1),
  211. (DataFrame([[1, 2], [2, 1]], columns=["A", "B"], index=["X", "Y"]), 0),
  212. ],
  213. )
  214. def test_background_gradient_gmap_wrong_dataframe(styler_blank, gmap, axis):
  215. # test giving a gmap in DataFrame but with wrong axis
  216. msg = "'gmap' is a DataFrame but underlying data for operations is a Series"
  217. with pytest.raises(ValueError, match=msg):
  218. styler_blank.background_gradient(gmap=gmap, axis=axis)._compute()
  219. def test_background_gradient_gmap_wrong_series(styler_blank):
  220. # test giving a gmap in Series form but with wrong axis
  221. msg = "'gmap' is a Series but underlying data for operations is a DataFrame"
  222. gmap = Series([1, 2], index=["X", "Y"])
  223. with pytest.raises(ValueError, match=msg):
  224. styler_blank.background_gradient(gmap=gmap, axis=None)._compute()
  225. def test_background_gradient_nullable_dtypes():
  226. # GH 50712
  227. df1 = DataFrame([[1], [0], [np.nan]], dtype=float)
  228. df2 = DataFrame([[1], [0], [None]], dtype="Int64")
  229. ctx1 = df1.style.background_gradient()._compute().ctx
  230. ctx2 = df2.style.background_gradient()._compute().ctx
  231. assert ctx1 == ctx2
  232. @pytest.mark.parametrize(
  233. "cmap",
  234. ["PuBu", mpl.colormaps["PuBu"]],
  235. )
  236. def test_bar_colormap(cmap):
  237. data = DataFrame([[1, 2], [3, 4]])
  238. ctx = data.style.bar(cmap=cmap, axis=None)._compute().ctx
  239. pubu_colors = {
  240. (0, 0): "#d0d1e6",
  241. (1, 0): "#056faf",
  242. (0, 1): "#73a9cf",
  243. (1, 1): "#023858",
  244. }
  245. for k, v in pubu_colors.items():
  246. assert v in ctx[k][1][1]
  247. def test_bar_color_raises(df):
  248. msg = "`color` must be string or list or tuple of 2 strings"
  249. with pytest.raises(ValueError, match=msg):
  250. df.style.bar(color={"a", "b"}).to_html()
  251. with pytest.raises(ValueError, match=msg):
  252. df.style.bar(color=["a", "b", "c"]).to_html()
  253. msg = "`color` and `cmap` cannot both be given"
  254. with pytest.raises(ValueError, match=msg):
  255. df.style.bar(color="something", cmap="something else").to_html()
  256. @pytest.mark.parametrize(
  257. "plot_method",
  258. ["scatter", "hexbin"],
  259. )
  260. def test_pass_colormap_instance(df, plot_method):
  261. # https://github.com/pandas-dev/pandas/issues/49374
  262. cmap = mpl.colors.ListedColormap([[1, 1, 1], [0, 0, 0]])
  263. df["c"] = df.A + df.B
  264. kwargs = dict(x="A", y="B", c="c", colormap=cmap)
  265. if plot_method == "hexbin":
  266. kwargs["C"] = kwargs.pop("c")
  267. getattr(df.plot, plot_method)(**kwargs)