Merge pull request #366 from eric-wieser/FancyArrowPatch

Add support for FancyArrowPatch objects

Author: nschloe

test/test_annotate_reference.tex

@@ -38,14 +38,14 @@
 4.6 -0.809016994374946
 4.8 0.30901699437495
 };
-\draw[->,red] (axis cs:4.5,1.5) -- (axis cs:4,1);
+\draw[->,draw=red] (axis cs:4.5,1.5) -- (axis cs:4,1);
 \draw (axis cs:4.5,1.5) node[
   scale=0.5,
   anchor=base west,
   text=black,
   rotate=0.0
 ]{text};
-\draw[->,black] (axis cs:0,1) ++(-50pt,30pt) -- (axis cs:0,1);
+\draw[->,draw=black] (axis cs:0,1) ++(-50pt,30pt) -- (axis cs:0,1);
 \draw (axis cs:0,1) ++(-50pt,30pt) node[
   scale=0.5,
   anchor=base west,

test/test_arrows_reference.tex

@@ -26,7 +26,7 @@
 \draw[draw=none,fill=color0] (axis cs:7.2,2.8) circle (0.2);
 \draw[draw=none,fill=color0] (axis cs:7.2,1.8) circle (0.2);
 \draw[draw=none,fill=color0] (axis cs:7.2,0.8) circle (0.2);
-\draw[-,black] (axis cs:2,6.8) -- (axis cs:3.2,6.8);
+\draw[-,draw=black] (axis cs:2,6.8) -- (axis cs:3.2,6.8);
 \draw (axis cs:2,6.8) node[
   scale=0.7,
   fill=white,
@@ -37,7 +37,7 @@
   text=black,
   rotate=0.0
 ]{-};
-\draw[->,black] (axis cs:2,5.8) -- (axis cs:3.2,5.8);
+\draw[->,draw=black] (axis cs:2,5.8) -- (axis cs:3.2,5.8);
 \draw (axis cs:2,5.8) node[
   scale=0.7,
   fill=white,
@@ -48,7 +48,7 @@
   text=black,
   rotate=0.0
 ]{-$>$};
-\draw[-|,black] (axis cs:2,4.8) -- (axis cs:3.2,4.8);
+\draw[-|,draw=black] (axis cs:2,4.8) -- (axis cs:3.2,4.8);
 \draw (axis cs:2,4.8) node[
   scale=0.7,
   fill=white,
@@ -59,7 +59,7 @@
   text=black,
   rotate=0.0
 ]{-[};
-\draw[-latex,black] (axis cs:2,3.8) -- (axis cs:3.2,3.8);
+\draw[-latex,draw=black] (axis cs:2,3.8) -- (axis cs:3.2,3.8);
 \draw (axis cs:2,3.8) node[
   scale=0.7,
   fill=white,
@@ -70,7 +70,7 @@
   text=black,
   rotate=0.0
 ]{-$|$$>$};
-\draw[<-,black] (axis cs:2,2.8) -- (axis cs:3.2,2.8);
+\draw[<-,draw=black] (axis cs:2,2.8) -- (axis cs:3.2,2.8);
 \draw (axis cs:2,2.8) node[
   scale=0.7,
   fill=white,
@@ -81,7 +81,7 @@
   text=black,
   rotate=0.0
 ]{$<$-};
-\draw[<->,black] (axis cs:2,1.8) -- (axis cs:3.2,1.8);
+\draw[<->,draw=black] (axis cs:2,1.8) -- (axis cs:3.2,1.8);
 \draw (axis cs:2,1.8) node[
   scale=0.7,
   fill=white,
@@ -92,7 +92,7 @@
   text=black,
   rotate=0.0
 ]{$<$-$>$};
-\draw[latex-,black] (axis cs:2,0.8) -- (axis cs:3.2,0.8);
+\draw[latex-,draw=black] (axis cs:2,0.8) -- (axis cs:3.2,0.8);
 \draw (axis cs:2,0.8) node[
   scale=0.7,
   fill=white,
@@ -103,7 +103,7 @@
   text=black,
   rotate=0.0
 ]{$<$$|$-};
-\draw[latex-latex,black] (axis cs:6,6.8) -- (axis cs:7.2,6.8);
+\draw[latex-latex,draw=black] (axis cs:6,6.8) -- (axis cs:7.2,6.8);
 \draw (axis cs:6,6.8) node[
   scale=0.7,
   fill=white,
@@ -114,7 +114,7 @@
   text=black,
   rotate=0.0
 ]{$<$$|$-$|$$>$};
-\draw[|-,black] (axis cs:6,5.8) -- (axis cs:7.2,5.8);
+\draw[|-,draw=black] (axis cs:6,5.8) -- (axis cs:7.2,5.8);
 \draw (axis cs:6,5.8) node[
   scale=0.7,
   fill=white,
@@ -125,7 +125,7 @@
   text=black,
   rotate=0.0
 ]{]-};
-\draw[|-|,black] (axis cs:6,4.8) -- (axis cs:7.2,4.8);
+\draw[|-|,draw=black] (axis cs:6,4.8) -- (axis cs:7.2,4.8);
 \draw (axis cs:6,4.8) node[
   scale=0.7,
   fill=white,
@@ -136,7 +136,7 @@
   text=black,
   rotate=0.0
 ]{]-[};
-\draw[-latex,very thick,black] (axis cs:6,3.8) -- (axis cs:7.2,3.8);
+\draw[-latex,very thick,draw=black] (axis cs:6,3.8) -- (axis cs:7.2,3.8);
 \draw (axis cs:6,3.8) node[
   scale=0.7,
   fill=white,
@@ -147,7 +147,7 @@
   text=black,
   rotate=0.0
 ]{fancy};
-\draw[-latex,very thick,black] (axis cs:6,2.8) -- (axis cs:7.2,2.8);
+\draw[-latex,very thick,draw=black] (axis cs:6,2.8) -- (axis cs:7.2,2.8);
 \draw (axis cs:6,2.8) node[
   scale=0.7,
   fill=white,
@@ -158,7 +158,7 @@
   text=black,
   rotate=0.0
 ]{simple};
-\draw[-latex,very thick,black] (axis cs:6,1.8) -- (axis cs:7.2,1.8);
+\draw[-latex,very thick,draw=black] (axis cs:6,1.8) -- (axis cs:7.2,1.8);
 \draw (axis cs:6,1.8) node[
   scale=0.7,
   fill=white,
@@ -169,7 +169,7 @@
   text=black,
   rotate=0.0
 ]{wedge};
-\draw[|-|,black] (axis cs:6,0.8) -- (axis cs:7.2,0.8);
+\draw[|-|,draw=black] (axis cs:6,0.8) -- (axis cs:7.2,0.8);
 \draw (axis cs:6,0.8) node[
   scale=0.7,
   fill=white,

test/test_patches.py

@@ -2,8 +2,16 @@
 
 
 def plot():
-    from matplotlib.patches import Circle, Ellipse, Polygon, Rectangle, Wedge
+    from matplotlib.patches import (
+        Circle,
+        Ellipse,
+        Polygon,
+        Rectangle,
+        Wedge,
+        FancyArrowPatch,
+    )
     from matplotlib.collections import PatchCollection
+    from matplotlib.path import Path
     from matplotlib import pyplot as plt
     import numpy as np
     import matplotlib as mpl
@@ -57,6 +65,18 @@ def plot():
     circle = Circle(xy=[0.0, 1.0], radius=0.5, color="r", alpha=0.4)
     ax.add_patch(circle)
 
+    arrow = FancyArrowPatch(posA=[0.25, 0.25], posB=[0.5, 0.25], arrowstyle="->")
+    ax.add_patch(arrow)
+
+    curved_arrow = FancyArrowPatch(
+        path=Path(
+            [(0.3, 0.3), (0.5, 1.0), (1.0, 0.8), (0.8, 0.3)],
+            [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4],
+        ),
+        arrowstyle="-|>",
+    )
+    ax.add_patch(curved_arrow)
+
     plt.colorbar(p)
 
     return fig

test/test_patches_reference.tex

@@ -198,6 +198,10 @@
 
 \draw[draw=none,fill=color0,fill opacity=0.4,rotate around={45:(axis cs:1,0.5)}] (axis cs:1,0.5) ellipse (0.5 and 0.25);
 \draw[draw=red,fill=red,opacity=0.4] (axis cs:0,1) circle (0.5);
+\draw[->,draw=black] (axis cs:0.25,0.25) -- (axis cs:0.5,0.25);
+\path [draw=black, -latex, draw=black]
+(axis cs:0.3,0.3)
+.. controls (axis cs:0.5,1) and (axis cs:1,0.8) .. (axis cs:0.8,0.3);
 \end{axis}
 
 \end{tikzpicture}

tikzplotlib/_patch.py

@@ -1,11 +1,25 @@
 import matplotlib as mpl
 
 from . import _path as mypath
+from ._text import _get_arrow_style
 
 
 def draw_patch(data, obj):
     """Return the PGFPlots code for patches.
     """
+    if isinstance(obj, mpl.patches.FancyArrowPatch):
+        data, draw_options = mypath.get_draw_options(
+            data,
+            obj,
+            obj.get_edgecolor(),
+            # get_fillcolor for the arrow refers to the head, not the path
+            None,
+            obj.get_linestyle(),
+            obj.get_linewidth(),
+            obj.get_hatch(),
+        )
+        return _draw_fancy_arrow(data, obj, draw_options)
+
     # Gather the draw options.
     data, draw_options = mypath.get_draw_options(
         data,
@@ -184,3 +198,19 @@ def _draw_circle(data, obj, draw_options):
     content += _patch_legend(obj, draw_options, "area legend")
 
     return data, content
+
+
+def _draw_fancy_arrow(data, obj, draw_options):
+    style = _get_arrow_style(obj, data)
+    ff = data["float format"]
+    if obj._posA_posB is not None:
+        posA, posB = obj._posA_posB
+        content = "\\draw[{{}}] (axis cs:{ff},{ff}) -- (axis cs:{ff},{ff});\n".format(
+            ff=ff
+        ).format(",".join(style), *posA, *posB)
+    else:
+        data, content, _, _ = mypath.draw_path(
+            data, obj._path_original, draw_options=draw_options + style
+        )
+    content += _patch_legend(obj, draw_options, "line legend")
+    return data, content

tikzplotlib/_text.py

@@ -1,4 +1,5 @@
 import matplotlib as mpl
+from matplotlib.patches import ArrowStyle
 
 from . import _color
 
@@ -164,6 +165,34 @@ def _parse_annotation_coords(ff, coords, xy):
         raise NotImplementedError
 
 
+def _get_arrow_style(obj, data):
+    # get a style string from a FancyArrowPatch
+    arrow_translate = {
+        ArrowStyle._style_list["-"]: ["-"],
+        ArrowStyle._style_list["->"]: ["->"],
+        ArrowStyle._style_list["<-"]: ["<-"],
+        ArrowStyle._style_list["<->"]: ["<->"],
+        ArrowStyle._style_list["|-|"]: ["|-|"],
+        ArrowStyle._style_list["-|>"]: ["-latex"],
+        ArrowStyle._style_list["<|-"]: ["latex-"],
+        ArrowStyle._style_list["<|-|>"]: ["latex-latex"],
+        ArrowStyle._style_list["]-["]: ["|-|"],
+        ArrowStyle._style_list["-["]: ["-|"],
+        ArrowStyle._style_list["]-"]: ["|-"],
+        ArrowStyle._style_list["fancy"]: ["-latex", "very thick"],
+        ArrowStyle._style_list["simple"]: ["-latex", "very thick"],
+        ArrowStyle._style_list["wedge"]: ["-latex", "very thick"],
+    }
+    style_cls = type(obj.get_arrowstyle())
+    try:
+        style = arrow_translate[style_cls]
+    except KeyError:
+        raise NotImplementedError("Unknown arrow style {}".format(style_cls))
+    else:
+        data, col, _ = _color.mpl_color2xcolor(data, obj.get_ec())
+        return style + ["draw=" + col]
+
+
 def _annotation(obj, data, content):
     ann_xy = obj.xy
     ann_xycoords = obj.xycoords
@@ -190,35 +219,9 @@ def _annotation(obj, data, content):
             # Anything else except for explicit positioning is not supported yet
             return obj.get_position()
 
-    # Create a basic tikz arrow
-    arrow_translate = {
-        "-": ["-"],
-        "->": ["->"],
-        "<-": ["<-"],
-        "<->": ["<->"],
-        "|-|": ["|-|"],
-        "-|>": ["-latex"],
-        "<|-": ["latex-"],
-        "<|-|>": ["latex-latex"],
-        "]-[": ["|-|"],
-        "-[": ["-|"],
-        "]-": ["|-"],
-        "fancy": ["-latex", "very thick"],
-        "simple": ["-latex", "very thick"],
-        "wedge": ["-latex", "very thick"],
-    }
-    arrow_style = []
-    if obj.arrowprops is not None:
-        if obj.arrowprops["arrowstyle"] is not None:
-            if obj.arrowprops["arrowstyle"] in arrow_translate:
-                arrow_style += arrow_translate[obj.arrowprops["arrowstyle"]]
-                data, col, _ = _color.mpl_color2xcolor(data, obj.arrow_patch.get_ec())
-                arrow_style.append(col)
-
-    if arrow_style:
-        the_arrow = ("\\draw[{}] {} -- {};\n").format(
-            ",".join(arrow_style), text_pos, xy_pos
-        )
+    if obj.arrow_patch:
+        style = ",".join(_get_arrow_style(obj.arrow_patch, data))
+        the_arrow = ("\\draw[{}] {} -- {};\n").format(style, text_pos, xy_pos)
         content.append(the_arrow)
     return text_pos