feat: accept Graphviz styling attributes in asdot (#139)

* docs: remove redundant Return and Return type sections
* feat: accept Graphviz attributes for edges and nodes
* feat: implement graphviz attributes for whole figure
* fix: remove redundant : from install page
* test: assert exact format of graphviz output

.cspell.json

@@ -119,6 +119,7 @@
         "astuple",
         "autoupdate",
         "bdist",
+        "bgcolor",
         "cano",
         "celltoolbar",
         "codacy",
@@ -198,6 +199,7 @@
         "startswith",
         "stm's",
         "sympify",
+        "sympy",
         "theano",
         "ticklabels",
         "tolist",

docs/install.md

@@ -12,7 +12,7 @@ The fastest way of installing this package is through PyPI or Conda:
 python3 -m pip install qrules
 ```
 
-::::
+:::
 
 :::{tabbed} Conda
 

docs/usage/visualize.ipynb

@@ -393,13 +393,59 @@
     "dot = qrules.io.asdot(reaction, collapse_graphs=True, render_node=False)\n",
     "graphviz.Source(dot)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Styling"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "<!-- cspell:ignore darkgreen fontcolor fontsize penwidth -->\n",
+    "The {func}`.asdot` function also takes [Graphviz attributes](https://graphviz.org/doc/info/attrs.html). These can be used to modify the layout of the whole figure. Examples are the [`size`](https://graphviz.org/docs/attrs/size), [`color`](https://graphviz.org/docs/attrs/color), and [`fontcolor`](https://graphviz.org/docs/attrs/fontcolor). Edges and nodes can be styled with `edge_style` and `node_style` respectively:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dot = qrules.io.asdot(\n",
+    "    reaction.transitions[0],\n",
+    "    render_node=True,\n",
+    "    size=12,\n",
+    "    bgcolor=\"transparent\",\n",
+    "    edge_style={\n",
+    "        \"color\": \"red\",\n",
+    "        \"arrowhead\": \"open\",\n",
+    "        \"fontcolor\": \"blue\",\n",
+    "        \"fontsize\": 25,\n",
+    "    },\n",
+    "    node_style={\n",
+    "        \"color\": \"gray\",\n",
+    "        \"penwidth\": 2,\n",
+    "        \"shape\": \"ellipse\",\n",
+    "        \"style\": \"dashed\",\n",
+    "    },\n",
+    ")\n",
+    "display(graphviz.Source(dot))\n"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
    "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.8.12"
   }
  },
  "nbformat": 4,

src/qrules/_system_control.py

@@ -2,7 +2,7 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, List, Optional, Set, Tuple, Type
+from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple, Type
 
 import attr
 
@@ -289,7 +289,8 @@ def __call__(
 
 
 def filter_graphs(
-    graphs: List[StateTransitionGraph], filters: List[Callable]
+    graphs: List[StateTransitionGraph],
+    filters: Iterable[Callable[[StateTransitionGraph], bool]],
 ) -> List[StateTransitionGraph]:
     r"""Implement filtering of a list of `.StateTransitionGraph` 's.
 
@@ -300,14 +301,6 @@ def filter_graphs(
     Note:
         For the more advanced user, lambda functions can be used as filters.
 
-    Args:
-        graphs ([`.StateTransitionGraph`]): list of graphs to be
-            filtered
-        filters (list): list of functions, which take a single
-            `.StateTransitionGraph` as an argument
-    Returns:
-        [`.StateTransitionGraph`]: filtered list of graphs
-
     Example:
         Selecting only the solutions, in which the :math:`\rho` decays via
         p-wave:

src/qrules/io/__init__.py

@@ -9,7 +9,7 @@
 import json
 from collections import abc
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import attr
 import yaml
@@ -91,6 +91,9 @@ def asdot(
     render_initial_state_id: bool = False,
     strip_spin: bool = False,
     collapse_graphs: bool = False,
+    edge_style: Optional[Dict[str, Any]] = None,
+    node_style: Optional[Dict[str, Any]] = None,
+    **figure_style: Any,
 ) -> str:
     """Convert a `object` to a DOT language `str`.
 
@@ -120,35 +123,51 @@ def asdot(
             See `.InteractionProperties` for more info.
 
         render_final_state_id: Add edge IDs for the final state edges.
-
         render_resonance_id: Add edge IDs for the intermediate state edges.
-
         render_initial_state_id: Add edge IDs for the initial state edges.
+        edge_style: Styling of a Graphviz edge.
+        node_style: Styling of a Graphviz node.
+        figure_style: Styling of the whole figure.
+
+    .. seealso::
+
+        See `Graphviz attributes <https://graphviz.org/doc/info/attrs.html>`_
+        for the available styling arguments.
 
     .. seealso:: :doc:`/usage/visualize`
     """
+    if edge_style is None:
+        edge_style = {}
+    if node_style is None:
+        node_style = {}
     if isinstance(instance, StateTransition):
         instance = instance.to_graph()
     if isinstance(instance, (ProblemSet, StateTransitionGraph, Topology)):
-        return _dot.graph_to_dot(
+        dot = _dot.graph_to_dot(
             instance,
             render_node=render_node,
             render_final_state_id=render_final_state_id,
             render_resonance_id=render_resonance_id,
             render_initial_state_id=render_initial_state_id,
+            edge_style=edge_style,
+            node_style=node_style,
         )
+        return _dot.insert_graphviz_styling(dot, graphviz_attrs=figure_style)
     if isinstance(instance, (ReactionInfo, StateTransitionCollection)):
         instance = instance.to_graphs()
     if isinstance(instance, abc.Iterable):
-        return _dot.graph_list_to_dot(
+        dot = _dot.graph_list_to_dot(
             instance,
             render_node=render_node,
             render_final_state_id=render_final_state_id,
             render_resonance_id=render_resonance_id,
             render_initial_state_id=render_initial_state_id,
             strip_spin=strip_spin,
             collapse_graphs=collapse_graphs,
+            edge_style=edge_style,
+            node_style=node_style,
         )
+        return _dot.insert_graphviz_styling(dot, graphviz_attrs=figure_style)
     raise NotImplementedError(
         f"Cannot convert a {instance.__class__.__name__} to DOT language"
     )