Allow edge style to be a list of styles for DiGraphs (#5131)

Support sequences of linestyles for directed edges in `draw_networkx_edges`.

2 files changed, 128 insertions, 2 deletions

diff --git a/networkx/drawing/nx_pylab.py b/networkx/drawing/nx_pylab.py
index e01b2054..836ff85e 100644
--- a/networkx/drawing/nx_pylab.py
+++ b/networkx/drawing/nx_pylab.py
@@ -540,9 +540,15 @@ def draw_networkx_edges(
         floats from 0-1. If numeric values are specified they will be
         mapped to colors using the edge_cmap and edge_vmin,edge_vmax parameters.
 
-    style : string (default=solid line)
+    style : string or array of strings (default='solid')
         Edge line style e.g.: '-', '--', '-.', ':'
         or words like 'solid' or 'dashed'.
+        Can be a single style or a sequence of styles with the same
+        length as the edge list.
+        If less styles than edges are given the styles will cycle.
+        If more styles than edges are given the styles will be used sequentially
+        and not be exhausted.
+        Also, `(offset, onoffseq)` tuples can be used as style instead of a strings.
         (See `matplotlib.patches.FancyArrowPatch`: `linestyle`)
 
     alpha : float or None (default=None)
@@ -822,6 +828,18 @@ def draw_networkx_edges(
             else:
                 line_width = width
 
+            if (
+                np.iterable(style)
+                and not isinstance(style, str)
+                and not isinstance(style, tuple)
+            ):
+                if len(style) == len(edge_pos):
+                    linestyle = style[i]
+                else:  # Cycle through styles
+                    linestyle = style[i % len(style)]
+            else:
+                linestyle = style
+
             arrow = mpl.patches.FancyArrowPatch(
                 (x1, y1),
                 (x2, y2),
@@ -832,7 +850,7 @@ def draw_networkx_edges(
                 color=arrow_color,
                 linewidth=line_width,
                 connectionstyle=_connectionstyle,
-                linestyle=style,
+                linestyle=linestyle,
                 zorder=1,
             )  # arrows go behind nodes
 
diff --git a/networkx/drawing/tests/test_pylab.py b/networkx/drawing/tests/test_pylab.py
index cb79708b..848650db 100644
--- a/networkx/drawing/tests/test_pylab.py
+++ b/networkx/drawing/tests/test_pylab.py
@@ -1,6 +1,7 @@
 """Unit tests for matplotlib drawing functions."""
 import os
 import itertools
+
 import pytest
 
 mpl = pytest.importorskip("matplotlib")
@@ -8,6 +9,7 @@ mpl.use("PS")
 plt = pytest.importorskip("matplotlib.pyplot")
 plt.rcParams["text.usetex"] = False
 
+
 import networkx as nx
 
 barbell = nx.barbell_graph(4, 6)
@@ -138,6 +140,112 @@ def test_edge_colors_and_widths():
         # plt.show()
 
 
+def test_linestyle():
+
+    np = pytest.importorskip("numpy")
+
+    def test_styles(edges, style="solid"):
+        """
+        Function to test the styles set for edges drawn as FancyArrowPatch(es)
+        TODO: It would be nice to run the same tests for LineCollection(s)
+        """
+
+        # we assume that if edges are not a LineCollection, they are drawn as FanceArrowPatches
+        if not isinstance(edges, mpl.collections.LineCollection):
+            for i, edge in enumerate(edges):
+                if isinstance(style, str) or isinstance(style, tuple):
+                    linestyle = style
+                elif np.iterable(style):
+                    if len(style) == len(edges):
+                        linestyle = style[i]
+                    else:  # Cycle through styles
+                        linestyle = style[i % len(style)]
+                else:
+                    linestyle = style
+                assert edge.get_linestyle() == linestyle
+
+    pos = nx.circular_layout(barbell)
+    for G in (barbell, barbell.to_directed()):
+
+        nx.draw_networkx_nodes(G, pos, node_color=[(1.0, 1.0, 0.2, 0.5)])
+        nx.draw_networkx_labels(G, pos)
+
+        # edge with default style
+        drawn_edges = nx.draw_networkx_edges(G, pos, edgelist=[(0, 1), (0, 2), (1, 2)])
+        test_styles(drawn_edges)
+
+        # global style
+
+        # edge with string style
+        style = "dashed"
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # edge with simplified string style
+        style = "--"
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # edge with tuple style
+        style = (1, (1, 1))
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # global style in list
+
+        # edge with string style in list
+        style = ["dashed"]
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # edge with simplified string style in list
+        style = ["--"]
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # edge with tuple style as list
+        style = [(1, (1, 1))]
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # styles for each edge
+
+        # edges with styles for each edge
+        style = ["--", "-", ":"]
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # edges with fewer styles than edges
+        style = ["--", "-"]
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # edges with more styles than edges
+        style = ["--", "-", ":", "-."]
+        drawn_edges = nx.draw_networkx_edges(
+            G, pos, edgelist=[(0, 1), (0, 2), (1, 2)], style=style
+        )
+        test_styles(drawn_edges, style=style)
+
+        # plt.show()
+
+
 def test_labels_and_colors():
     G = nx.cubical_graph()
     pos = nx.spring_layout(G)  # positions for all nodes