Repository: networkx/networkx
Pull request: Add option for arrowsize to be a list
This PR fixes #3272 

Sample code:
```Python
>>> import networkx as nx
>>> import matplotlib.pyplot as plt
>>> G = nx.DiGraph([(0, 1), (0, 2), (1, 3)])
>>> arrowsize=[10, 20, 30]
>>> nx.draw(G, arrowsize=arrowsize)
>>> plt.show()
```
Produces:
![image](https://user-images.githubusercontent.com/33282622/138859013-0e4fb354-c766-4097-88c8-1f33429d5714.png)


Problem statement:
Allow for different arrow sizes
In directed graph drawing, it's much clearer to tell edge weights with different arrow sizes.

In nx_pylab.py, line 676, arrowsize can easily be input as a list if just a few lines are added:

`if len(arrowsize) > 1:
    mutation_scale = arrowsize[i]
else:
    mutation_scale = arrowsize[0]`

Hints from issue discussion:

Commit message: Add option for arrowsize to be a list

Patch to review:
diff --git a/networkx/drawing/nx_pylab.py b/networkx/drawing/nx_pylab.py
index 836ff85ea42..b6073adb242 100644
--- a/networkx/drawing/nx_pylab.py
+++ b/networkx/drawing/nx_pylab.py
@@ -155,10 +155,12 @@ def draw_networkx(G, pos=None, arrows=None, with_labels=True, **kwds):
         For directed graphs, choose the style of the arrowsheads.
         See `matplotlib.patches.ArrowStyle` for more options.
 
-    arrowsize : int (default=10)
+    arrowsize : int or list (default=10)
         For directed graphs, choose the size of the arrow head's length and
-        width. See `matplotlib.patches.FancyArrowPatch` for attribute
-        `mutation_scale` for more info.
+        width.
+        If `list`, assign different size on each arrow head's length and width.
+        See `matplotlib.patches.FancyArrowPatch` for attribute `mutation_scale`
+        for more info.
 
     with_labels :  bool (default=True)
         Set to True to draw labels on the nodes.
@@ -750,7 +752,12 @@ def to_marker_edge(marker_size, marker):
 
         # Draw arrows with `matplotlib.patches.FancyarrowPatch`
         arrow_collection = []
-        mutation_scale = arrowsize  # scale factor of arrow head
+
+        if isinstance(arrowsize, list):
+            if len(arrowsize) != len(edge_pos):
+                raise ValueError("arrowsize should have the same length as G.edges")
+        else:
+            mutation_scale = arrowsize  # scale factor of arrow head
 
         base_connection_style = mpl.patches.ConnectionStyle(connectionstyle)
 
@@ -798,6 +805,11 @@ def _connectionstyle(posA, posB, *args, **kwargs):
             x2, y2 = dst
             shrink_source = 0  # space from source to tail
             shrink_target = 0  # space from  head to target
+
+            if isinstance(arrowsize, list):
+                # Scale each factor of each arrow based on arrowsize list
+                mutation_scale = arrowsize[i]
+
             if np.iterable(node_size):  # many node sizes
                 source, target = edgelist[i][:2]
                 source_node_size = node_size[nodelist.index(source)]
diff --git a/networkx/drawing/tests/test_pylab.py b/networkx/drawing/tests/test_pylab.py
index f4a1a73e64c..88334e06a01 100644
--- a/networkx/drawing/tests/test_pylab.py
+++ b/networkx/drawing/tests/test_pylab.py
@@ -478,6 +478,23 @@ def test_error_invalid_kwds():
         nx.draw(barbell, foo="bar")
 
 
+def test_draw_networkx_arrowsize_incorrect_size():
+    G = nx.DiGraph([(0, 1), (0, 2), (0, 3), (1, 3)])
+    arrowsize = [1, 2, 3]
+    with pytest.raises(
+        ValueError, match="arrowsize should have the same length as G.edges"
+    ):
+        nx.draw(G, arrowsize=arrowsize)
+
+
+@pytest.mark.parametrize("arrowsize", (30, [10, 20, 30]))
+def test_draw_edges_arrowsize(arrowsize):
+    G = nx.DiGraph([(0, 1), (0, 2), (1, 2)])
+    pos = {0: (0, 0), 1: (0, 1), 2: (1, 0)}
+    edges = nx.draw_networkx_edges(G, pos=pos, arrowsize=arrowsize)
+    assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
+
+
 def test_np_edgelist():
     # see issue #4129
     np = pytest.importorskip("numpy")


Code review comments:
diff_hunk: @@ -478,6 +478,23 @@ def test_error_invalid_kwds():
         nx.draw(barbell, foo="bar")
 
 
+def test_draw_networkx_arrowsize_incorrect_size():
+    G = nx.DiGraph([(0, 1), (0, 2), (0, 3), (1, 3)])
+    arrowsize = [1, 2, 3]
+    with pytest.raises(
+        ValueError, match="arrowsize should have the same length as G.edges"
+    ):
+        nx.draw(G, arrowsize=arrowsize)
+
+
+@pytest.mark.parametrize("arrowsize", (30, [10, 20, 30]))
+def test_draw_edges_arrowsize(arrowsize):
+    G = nx.DiGraph([(0, 1), (0, 2), (1, 2)])
+    pos = {0: (0, 0), 1: (0, 1), 2: (1, 0)}
+    edges = nx.draw_networkx_edges(G, pos=pos, arrowsize=arrowsize)
+    assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
path: networkx/drawing/tests/test_pylab.py
text: @user1:
Thanks for the tests, these are great!

I think we can even go one step further and check the `mutation_scale` setting even more explicitly, something like:

```python
    for fap, expected in zip(edges, arrowsize):
        assert fap.get_mutation_scale() == expected
```

Of course, that messes up the parametrization a bit because `30` isn't iterable, so you'd either have to break it up into two tests (say `test_draw_edges_arrowsize_int` and `test_draw_edges_arrowsize_list`) or add something like `arrowsize = itertools.repeat(arrowsize) if isinstance(arrowsize, int) else arrowsize` before the for-loop.

@author:
I really like the second suggestion using `itertools`! Thanks

diff_hunk: @@ -750,7 +752,12 @@ def to_marker_edge(marker_size, marker):
 
         # Draw arrows with `matplotlib.patches.FancyarrowPatch`
         arrow_collection = []
-        mutation_scale = arrowsize  # scale factor of arrow head
+
+        if isinstance(arrowsize, list):
+            if len(arrowsize) != len(edge_pos):
+                raise ValueError("arrowsize should have the same length as G.edges")
path: networkx/drawing/nx_pylab.py
text: @user1:
This error message will be incorrect in the case where a user passes in an explicit `edgelist` that has fewer edges than `G.edges`. A simple fix might be to simply replace "G.edges" with "edgelist", which is slightly more accurate (the default edgelist is `list(G.edges)`) if not quite perfect!

@author:
Yeah you are right I will change the message