add edge_label_kwargs interface to :meth:.PipelineState.plot_graph

Author: Aethor

renard/pipeline/core.py

@@ -428,6 +428,7 @@ def plot_graph(
         node_kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
         edge_kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
         label_kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
+        edge_label_kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
         tight_layout: bool = False,
         legend: bool = False,
     ):
@@ -458,7 +459,10 @@ def plot_graph(
         :param node_kwargs: passed to :func:`nx.draw_networkx_nodes`
         :param edge_kwargs: passed to :func:`nx.draw_networkx_nodes`
         :param label_kwargs: passed to :func:`nx.draw_networkx_labels`
-        :param tight_layout: if ``True``, will use matplotlib's tight_layout
+        :param edge_label_kwargs: passed to
+            :func:`nx.draw_networkx_labels`
+        :param tight_layout: if ``True``, will use matplotlib's
+            tight_layout
         :param legend: passed to :func:`.plot_nx_graph_reasonably`
         """
         import matplotlib.pyplot as plt
@@ -481,6 +485,7 @@ def plot_graph(
             assert not isinstance(node_kwargs, list)
             assert not isinstance(edge_kwargs, list)
             assert not isinstance(label_kwargs, list)
+            assert not isinstance(edge_label_kwargs, list)
             if tight_layout:
                 fig.tight_layout()
             plot_nx_graph_reasonably(
@@ -490,6 +495,7 @@ def plot_graph(
                 node_kwargs=node_kwargs,
                 edge_kwargs=edge_kwargs,
                 label_kwargs=label_kwargs,
+                edge_label_kwargs=edge_label_kwargs,
                 legend=legend,
             )
             return
@@ -504,6 +510,10 @@ def plot_graph(
         assert isinstance(edge_kwargs, list)
         label_kwargs = label_kwargs or [{} for _ in range(len(self.character_network))]
         assert isinstance(label_kwargs, list)
+        edge_label_kwargs = edge_label_kwargs or [
+            {} for _ in range(len(self.character_network))
+        ]
+        assert isinstance(edge_label_kwargs, list)
 
         if fig is None:
             fig, ax = plt.subplots()
@@ -541,6 +551,7 @@ def update(slider_value):
                 node_kwargs=node_kwargs[slider_i],
                 edge_kwargs=edge_kwargs[slider_i],
                 label_kwargs=label_kwargs[slider_i],
+                edge_label_kwargs=edge_label_kwargs[slider_i],
                 legend=legend,
             )
             ax.set_xlim(-1.2, 1.2)

renard/plot_utils.py

@@ -25,7 +25,7 @@ def plot_nx_graph_reasonably(
     node_kwargs: Optional[Dict[str, Any]] = None,
     edge_kwargs: Optional[Dict[str, Any]] = None,
     label_kwargs: Optional[Dict[str, Any]] = None,
-    edge_labels_kwargs: Optional[Dict[str, Any]] = None,
+    edge_label_kwargs: Optional[Dict[str, Any]] = None,
     legend: bool = False,
 ):
     """Try to plot a :class:`nx.Graph` with 'reasonable' parameters
@@ -82,17 +82,17 @@ def plot_nx_graph_reasonably(
     edge_kwargs["alpha"] = edge_kwargs.get("alpha", 0.35)
     nx.draw_networkx_edges(G, pos, ax=ax, **edge_kwargs)
 
-    edge_labels_kwargs = edge_labels_kwargs or {}
-    if not "edge_labels" in edge_labels_kwargs:
-        edge_labels_kwargs["edge_labels"] = {
+    edge_label_kwargs = edge_label_kwargs or {}
+    if not "edge_labels" in edge_label_kwargs:
+        edge_label_kwargs["edge_labels"] = {
             (char1, char2): data.get("relations", "")
             for char1, char2, data in G.edges.data()
         }
-        for (char1, char2), rel in edge_labels_kwargs["edge_labels"].items():
+        for (char1, char2), rel in edge_label_kwargs["edge_labels"].items():
             if rel == set():
-                edge_labels_kwargs["edge_labels"][(char1, char2)] = ""
-    edge_labels_kwargs["font_size"] = edge_labels_kwargs.get("font_size", 6)
-    nx.draw_networkx_edge_labels(G, pos, ax=ax, **edge_labels_kwargs)
+                edge_label_kwargs["edge_labels"][(char1, char2)] = ""
+    edge_label_kwargs["font_size"] = edge_label_kwargs.get("font_size", 6)
+    nx.draw_networkx_edge_labels(G, pos, ax=ax, **edge_label_kwargs)
 
     label_kwargs = label_kwargs or {}
     label_kwargs["verticalalignment"] = label_kwargs.get("verticalalignment", "top")