test_style.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. import pytest
  2. from pandas import Series
  3. pytest.importorskip("matplotlib")
  4. from pandas.plotting._matplotlib.style import get_standard_colors
  5. class TestGetStandardColors:
  6. @pytest.mark.parametrize(
  7. "num_colors, expected",
  8. [
  9. (3, ["red", "green", "blue"]),
  10. (5, ["red", "green", "blue", "red", "green"]),
  11. (7, ["red", "green", "blue", "red", "green", "blue", "red"]),
  12. (2, ["red", "green"]),
  13. (1, ["red"]),
  14. ],
  15. )
  16. def test_default_colors_named_from_prop_cycle(self, num_colors, expected):
  17. import matplotlib as mpl
  18. from matplotlib.pyplot import cycler
  19. mpl_params = {
  20. "axes.prop_cycle": cycler(color=["red", "green", "blue"]),
  21. }
  22. with mpl.rc_context(rc=mpl_params):
  23. result = get_standard_colors(num_colors=num_colors)
  24. assert result == expected
  25. @pytest.mark.parametrize(
  26. "num_colors, expected",
  27. [
  28. (1, ["b"]),
  29. (3, ["b", "g", "r"]),
  30. (4, ["b", "g", "r", "y"]),
  31. (5, ["b", "g", "r", "y", "b"]),
  32. (7, ["b", "g", "r", "y", "b", "g", "r"]),
  33. ],
  34. )
  35. def test_default_colors_named_from_prop_cycle_string(self, num_colors, expected):
  36. import matplotlib as mpl
  37. from matplotlib.pyplot import cycler
  38. mpl_params = {
  39. "axes.prop_cycle": cycler(color="bgry"),
  40. }
  41. with mpl.rc_context(rc=mpl_params):
  42. result = get_standard_colors(num_colors=num_colors)
  43. assert result == expected
  44. @pytest.mark.parametrize(
  45. "num_colors, expected_name",
  46. [
  47. (1, ["C0"]),
  48. (3, ["C0", "C1", "C2"]),
  49. (
  50. 12,
  51. [
  52. "C0",
  53. "C1",
  54. "C2",
  55. "C3",
  56. "C4",
  57. "C5",
  58. "C6",
  59. "C7",
  60. "C8",
  61. "C9",
  62. "C0",
  63. "C1",
  64. ],
  65. ),
  66. ],
  67. )
  68. def test_default_colors_named_undefined_prop_cycle(self, num_colors, expected_name):
  69. import matplotlib as mpl
  70. import matplotlib.colors as mcolors
  71. with mpl.rc_context(rc={}):
  72. expected = [mcolors.to_hex(x) for x in expected_name]
  73. result = get_standard_colors(num_colors=num_colors)
  74. assert result == expected
  75. @pytest.mark.parametrize(
  76. "num_colors, expected",
  77. [
  78. (1, ["red", "green", (0.1, 0.2, 0.3)]),
  79. (2, ["red", "green", (0.1, 0.2, 0.3)]),
  80. (3, ["red", "green", (0.1, 0.2, 0.3)]),
  81. (4, ["red", "green", (0.1, 0.2, 0.3), "red"]),
  82. ],
  83. )
  84. def test_user_input_color_sequence(self, num_colors, expected):
  85. color = ["red", "green", (0.1, 0.2, 0.3)]
  86. result = get_standard_colors(color=color, num_colors=num_colors)
  87. assert result == expected
  88. @pytest.mark.parametrize(
  89. "num_colors, expected",
  90. [
  91. (1, ["r", "g", "b", "k"]),
  92. (2, ["r", "g", "b", "k"]),
  93. (3, ["r", "g", "b", "k"]),
  94. (4, ["r", "g", "b", "k"]),
  95. (5, ["r", "g", "b", "k", "r"]),
  96. (6, ["r", "g", "b", "k", "r", "g"]),
  97. ],
  98. )
  99. def test_user_input_color_string(self, num_colors, expected):
  100. color = "rgbk"
  101. result = get_standard_colors(color=color, num_colors=num_colors)
  102. assert result == expected
  103. @pytest.mark.parametrize(
  104. "num_colors, expected",
  105. [
  106. (1, [(0.1, 0.2, 0.3)]),
  107. (2, [(0.1, 0.2, 0.3), (0.1, 0.2, 0.3)]),
  108. (3, [(0.1, 0.2, 0.3), (0.1, 0.2, 0.3), (0.1, 0.2, 0.3)]),
  109. ],
  110. )
  111. def test_user_input_color_floats(self, num_colors, expected):
  112. color = (0.1, 0.2, 0.3)
  113. result = get_standard_colors(color=color, num_colors=num_colors)
  114. assert result == expected
  115. @pytest.mark.parametrize(
  116. "color, num_colors, expected",
  117. [
  118. ("Crimson", 1, ["Crimson"]),
  119. ("DodgerBlue", 2, ["DodgerBlue", "DodgerBlue"]),
  120. ("firebrick", 3, ["firebrick", "firebrick", "firebrick"]),
  121. ],
  122. )
  123. def test_user_input_named_color_string(self, color, num_colors, expected):
  124. result = get_standard_colors(color=color, num_colors=num_colors)
  125. assert result == expected
  126. @pytest.mark.parametrize("color", ["", [], (), Series([], dtype="object")])
  127. def test_empty_color_raises(self, color):
  128. with pytest.raises(ValueError, match="Invalid color argument"):
  129. get_standard_colors(color=color, num_colors=1)
  130. @pytest.mark.parametrize(
  131. "color",
  132. [
  133. "bad_color",
  134. ("red", "green", "bad_color"),
  135. (0.1,),
  136. (0.1, 0.2),
  137. (0.1, 0.2, 0.3, 0.4, 0.5), # must be either 3 or 4 floats
  138. ],
  139. )
  140. def test_bad_color_raises(self, color):
  141. with pytest.raises(ValueError, match="Invalid color"):
  142. get_standard_colors(color=color, num_colors=5)