Add test for axisline_style

oscargus · oscargus · commit 3aedfff5e797 · 2022-11-24T18:12:43.000+01:00
diff --git a/lib/mpl_toolkits/axisartist/axisline_style.py b/lib/mpl_toolkits/axisartist/axisline_style.py
@@ -2,9 +2,10 @@
 
 import numpy as np
 
+import matplotlib as mpl
 from matplotlib.patches import _Style, FancyArrowPatch
-from matplotlib.transforms import IdentityTransform
 from matplotlib.path import Path
+from matplotlib.transforms import IdentityTransform
 
 
 class _FancyAxislineStyle:
@@ -54,10 +55,10 @@ def set_path(self, path):
         def draw(self, renderer):
             """
             Draw the axis line.
-             1) transform the path to the display coordinate.
-             2) extend the path to make a room for arrow
-             3) update the path of the FancyArrowPatch.
-             4) draw
+             1) Transform the path to the display coordinate.
+             2) Extend the path to make a room for arrow.
+             3) Update the path of the FancyArrowPatch.
+             4) Draw.
             """
             path_in_disp = self._line_transform.transform_path(self._line_path)
             mutation_size = self.get_mutation_scale()  # line_mutation_scale()
@@ -67,9 +68,15 @@ def draw(self, renderer):
             FancyArrowPatch.draw(self, renderer)
 
     class FilledArrow(SimpleArrow):
-        """The artist class that will be returned for SimpleArrow style."""
+        """The artist class that will be returned for FilledArrow style."""
         _ARROW_STYLE = "-|>"
 
+        def __init__(self, axis_artist, line_path, transform,
+                     line_mutation_scale, facecolor):
+            super().__init__(axis_artist, line_path, transform,
+                             line_mutation_scale)
+            self.set_facecolor(facecolor)
+
 
 class AxislineStyle(_Style):
     """
@@ -131,7 +138,6 @@ def __init__(self, size=1):
             super().__init__()
 
         def new_line(self, axis_artist, transform):
-
             linepath = Path([(0, 0), (0, 1)])
             axisline = self.ArrowAxisClass(axis_artist, linepath, transform,
                                            line_mutation_scale=self.size)
@@ -140,6 +146,27 @@ def new_line(self, axis_artist, transform):
     _style_list["->"] = SimpleArrow
 
     class FilledArrow(SimpleArrow):
+        def __init__(self, size=1, facecolor=mpl.rcParams['axes.edgecolor']):
+            """
+            Parameters
+            ----------
+            size : float
+                Size of the arrow as a fraction of the ticklabel size.
+            facecolor : color, default: :rc:`axes.edgecolor`
+                Fill color.
+            """
+
+            self.size = size
+            self._facecolor = facecolor
+            super().__init__(size=size)
+
+        def new_line(self, axis_artist, transform):
+            linepath = Path([(0, 0), (0, 1)])
+            axisline = self.ArrowAxisClass(axis_artist, linepath, transform,
+                                           line_mutation_scale=self.size,
+                                           facecolor=self._facecolor)
+            return axisline
+
         ArrowAxisClass = _FancyAxislineStyle.FilledArrow
 
     _style_list["-|>"] = FilledArrow
diff --git a/lib/mpl_toolkits/axisartist/tests/test_axislines.py b/lib/mpl_toolkits/axisartist/tests/test_axislines.py
@@ -3,7 +3,7 @@
 from matplotlib.testing.decorators import image_comparison
 from matplotlib.transforms import IdentityTransform
 
-from mpl_toolkits.axisartist.axislines import SubplotZero, Subplot
+from mpl_toolkits.axisartist.axislines import AxesZero, SubplotZero, Subplot
 from mpl_toolkits.axisartist import Axes, SubplotHost
 
 
@@ -90,3 +90,30 @@ def test_ParasiteAxesAuxTrans():
         ax1.set_ylim((0, 5))
 
     ax2.contour(xx, yy, data, colors='k')
+
+
+@image_comparison(['axisline_style.png'], remove_text=True, style='mpl20')
+def test_axisline_style():
+    fig = plt.figure(figsize=(2, 2))
+    ax = fig.add_subplot(axes_class=AxesZero)
+    ax.axis["xzero"].set_axisline_style("-|>")
+    ax.axis["xzero"].set_visible(True)
+    ax.axis["yzero"].set_axisline_style("->")
+    ax.axis["yzero"].set_visible(True)
+
+    for direction in ("left", "right", "bottom", "top"):
+        ax.axis[direction].set_visible(False)
+
+
+@image_comparison(['axisline_style_size_color.png'], remove_text=True,
+                  style='mpl20')
+def test_axisline_style_size_color():
+    fig = plt.figure(figsize=(2, 2))
+    ax = fig.add_subplot(axes_class=AxesZero)
+    ax.axis["xzero"].set_axisline_style("-|>", size=2.0, facecolor='r')
+    ax.axis["xzero"].set_visible(True)
+    ax.axis["yzero"].set_axisline_style("->, size=1.5")
+    ax.axis["yzero"].set_visible(True)
+
+    for direction in ("left", "right", "bottom", "top"):
+        ax.axis[direction].set_visible(False)
