Edge attributes support

diagrams/__init__.py

@@ -3,7 +3,7 @@
 from hashlib import md5
 from pathlib import Path
 from random import getrandbits
-from typing import List, Union
+from typing import List, Union, Dict
 
 from graphviz import Digraph
 
@@ -75,15 +75,15 @@ class Diagram:
     # TODO: Label position option
     # TODO: Save directory option (filename + directory?)
     def __init__(
-        self,
-        name: str = "",
-        filename: str = "",
-        direction: str = "LR",
-        outformat: str = "png",
-        show: bool = True,
-        graph_attr: dict = {},
-        node_attr: dict = {},
-        edge_attr: dict = {},
+            self,
+            name: str = "",
+            filename: str = "",
+            direction: str = "LR",
+            outformat: str = "png",
+            show: bool = True,
+            graph_attr: dict = {},
+            node_attr: dict = {},
+            edge_attr: dict = {},
     ):
         """Diagram represents a global diagrams context.
 
@@ -129,6 +129,9 @@ def __init__(
 
         self.show = show
 
+    def __str__(self) -> str:
+        return str(self.dot)
+
     def __enter__(self):
         setdiagram(self)
         return self
@@ -160,15 +163,9 @@ def node(self, hashid: str, label: str, **attrs) -> None:
         """Create a new node."""
         self.dot.node(hashid, label=label, **attrs)
 
-    def connect(self, node: "Node", node2: "Node", directed=True) -> None:
+    def connect(self, node: "Node", node2: "Node", edge: "Edge") -> None:
         """Connect the two Nodes."""
-        attrs = {"dir": "none"} if not directed else {}
-        self.dot.edge(node.hashid, node2.hashid, **attrs)
-
-    def reverse(self, node: "Node", node2: "Node", directed=True) -> None:
-        """Connect the two Nodes in reverse direction."""
-        attrs = {"dir": "none"} if not directed else {"dir": "back"}
-        self.dot.edge(node.hashid, node2.hashid, **attrs)
+        self.dot.edge(node.hashid, node2.hashid, **edge.to_dot())
 
     def subgraph(self, dot: Digraph) -> None:
         """Create a subgraph for clustering"""
@@ -302,85 +299,90 @@ def __repr__(self):
         _name = self.__class__.__name__
         return f"<{self._provider}.{self._type}.{_name}>"
 
-    def __sub__(self, other: Union["Node", List["Node"]]):
-        """Implement Self - Node and Self - [Nodes]"""
-        if not isinstance(other, list):
-            return self.connect(other, directed=False)
-        for node in other:
-            self.connect(node, directed=False)
-        return other
-
-    def __rsub__(self, other: List["Node"]):
-        """
-        Called for [Nodes] - Self because list of Nodes don't have
-        __sub__ operators.
-        """
-        self.__sub__(other)
+    def __sub__(self, other: Union["Node", List["Node"], "Edge"]):
+        """Implement Self - Node, Self - [Nodes] and Self - Edge."""
+        if isinstance(other, list):
+            for node in other:
+                self.connect(node, Edge(self))
+            return other
+        elif isinstance(other, Node):
+            return self.connect(other, Edge(self))
+        else:
+            other.node = self
+            return other
+
+    def __rsub__(self, other: Union[List["Node"], List["Edge"]]):
+        """ Called for [Nodes] and [Edges] - Self because list don't have __sub__ operators. """
+        for o in other:
+            if isinstance(o, Edge):
+                o.connect(self)
+            else:
+                o.connect(self, Edge(self))
         return self
 
-    def __rshift__(self, other: Union["Node", List["Node"]]):
-        """Implements Self >> Node and Self >> [Nodes]."""
-        if not isinstance(other, list):
-            return self.connect(other)
-        for node in other:
-            self.connect(node)
-        return other
-
-    def __lshift__(self, other: Union["Node", List["Node"]]):
-        """Implements Self << Node and Self << [Nodes]."""
-        if not isinstance(other, list):
-            return self.reverse(other)
-        for node in other:
-            self.reverse(node)
-        return other
-
-    def __rrshift__(self, other: List["Node"]):
-        """
-        Called for [Nodes] >> Self because list of Nodes don't have
-        __rshift__ operators.
-        """
-        for node in other:
-            node.connect(self)
+    def __rshift__(self, other: Union["Node", List["Node"], "Edge"]):
+        """Implements Self >> Node, Self >> [Nodes] and Self Edge."""
+        if isinstance(other, list):
+            for node in other:
+                self.connect(node, Edge(self, forward=True))
+            return other
+        elif isinstance(other, Node):
+            return self.connect(other, Edge(self, forward=True))
+        else:
+            other.forward = True
+            other.node = self
+            return other
+
+    def __lshift__(self, other: Union["Node", List["Node"], "Edge"]):
+        """Implements Self << Node, Self << [Nodes] and Self << Edge."""
+        if isinstance(other, list):
+            for node in other:
+                self.connect(node, Edge(self, reverse=True))
+            return other
+        elif isinstance(other, Node):
+            return self.connect(other, Edge(self, reverse=True))
+        else:
+            other.reverse = True
+            return other.connect(self)
+
+    def __rrshift__(self, other: Union[List["Node"], List["Edge"]]):
+        """Called for [Nodes] and [Edges] >> Self because list don't have __rshift__ operators."""
+        for o in other:
+            if isinstance(o, Edge):
+                o.forward = True
+                o.connect(self)
+            else:
+                o.connect(self, Edge(self, forward=True))
         return self
 
-    def __rlshift__(self, other: List["Node"]):
-        """
-        Called for [Nodes] << Self because list of Nodes don't have
-        __lshift__ operators.
-        """
-        for node in other:
-            node.reverse(self)
+    def __rlshift__(self, other: Union[List["Node"], List["Edge"]]):
+        """Called for [Nodes] << Self because list of Nodes don't have __lshift__ operators."""
+        for o in other:
+            if isinstance(o, Edge):
+                o.reverse = True
+                o.connect(self)
+            else:
+                o.connect(self, Edge(self, reverse=True))
         return self
 
     @property
     def hashid(self):
         return self._hash
 
     # TODO: option for adding flow description to the connection edge
-    def connect(self, node: "Node", directed=True):
+    def connect(self, node: "Node", edge: "Edge"):
         """Connect to other node.
 
         :param node: Other node instance.
-        :param directed: Whether the flow is directed or not.
-        :return: Connected node.
-        """
-        if not isinstance(node, Node):
-            ValueError(f"{node} is not a valid Node")
-        # An edge must be added on the global diagrams, not a cluster.
-        self._diagram.connect(self, node, directed)
-        return node
-
-    def reverse(self, node: "Node", directed=True):
-        """Connect to other node in reverse direction.
-
-        :param node: Other node instance.
-        :param directed: Whether the flow is directed or not.
+        :param edge: Type of the edge.
         :return: Connected node.
         """
         if not isinstance(node, Node):
             ValueError(f"{node} is not a valid Node")
+        if not isinstance(node, Edge):
+            ValueError(f"{node} is not a valid Edge")
         # An edge must be added on the global diagrams, not a cluster.
-        self._diagram.reverse(self, node, directed)
+        self._diagram.connect(self, node, edge)
         return node
 
     @staticmethod
@@ -392,4 +394,111 @@ def _load_icon(self):
         return os.path.join(basedir.parent, self._icon_dir, self._icon)
 
 
+class Edge:
+    """Edge represents an edge between two nodes."""
+
+    def __init__(self,
+                 node: "Node" = None,
+                 forward: bool = False,
+                 reverse: bool = False,
+                 label: str = "",
+                 color: str = "",
+                 style: str = ""
+                 ):
+        """Edge represents an edge between two nodes.
+
+        :param node: Parent node.
+        :param forward: Points forward.
+        :param reverse: Points backward.
+        :param label: Edge label.
+        :param color: Edge color.
+        :param style: Edge style.
+        """
+        if node is not None:
+            assert isinstance(node, Node)
+
+        self.node = node
+        self.forward = forward
+        self.reverse = reverse
+        self.label = label
+        self.color = color
+        self.style = style
+
+    def __sub__(self, other: Union["Node", "Edge", List["Node"]]):
+        """Implement Self - Node or Edge and Self - [Nodes]"""
+        return self.connect(other)
+
+    def __rsub__(self, other: Union[List["Node"], List["Edge"]]) -> List["Edge"]:
+        """Called for [Nodes] or [Edges] - Self because list don't have __sub__ operators."""
+        return self.append(other)
+
+    def __rshift__(self, other: Union["Node", "Edge", List["Node"]]):
+        """Implements Self >> Node or Edge and Self >> [Nodes]."""
+        self.forward = True
+        return self.connect(other)
+
+    def __lshift__(self, other: Union["Node", "Edge", List["Node"]]):
+        """Implements Self << Node or Edge and Self << [Nodes]."""
+        self.reverse = True
+        return self.connect(other)
+
+    def __rrshift__(self, other: Union[List["Node"], List["Edge"]]) -> List["Edge"]:
+        """Called for [Nodes] or [Edges] >> Self because list of Edges don't have __rshift__ operators."""
+        return self.append(other, forward=True)
+
+    def __rlshift__(self, other: Union[List["Node"], List["Edge"]]) -> List["Edge"]:
+        """Called for [Nodes] or [Edges] << Self because list of Edges don't have __lshift__ operators."""
+        return self.append(other, reverse=True)
+
+    def append(self, other: Union[List["Node"], List["Edge"]], forward=None, reverse=None) -> List["Edge"]:
+        result = []
+        for o in other:
+            if isinstance(o, Edge):
+                o.color = self.color
+                o.label = self.label
+                o.style = self.style
+                o.forward = forward if forward is not None else o.forward
+                o.reverse = forward if forward is not None else o.reverse
+                result.append(o)
+            else:
+                result.append(Edge(o, forward=forward, reverse=reverse,
+                                   color=self.color, label=self.label, style=self.style))
+        return result
+
+    def connect(self, other: Union["Node", "Edge", List["Node"]]):
+        if isinstance(other, list):
+            for node in other:
+                self.node.connect(node, self)
+            return other
+        elif isinstance(other, Edge):
+            self.label = other.label
+            self.color = other.color
+            self.style = other.style
+            return self
+        else:
+            if self.node is not None:
+                return self.node.connect(other, self)
+            else:
+                self.node = other
+                return self
+
+    def to_dot(self) -> Dict:
+        dot = {}
+        if self.forward and self.reverse:
+            dot['dir'] = 'both'
+        elif self.forward:
+            dot['dir'] = 'forward'
+        elif self.reverse:
+            dot['dir'] = 'back'
+        else:
+            dot['dir'] = 'none'
+        if self.label:
+            dot['label'] = self.label
+        if self.color:
+            dot['color'] = self.color
+        if self.style:
+            dot['style'] = self.style
+        return dot
+
+
 Group = Cluster

docs/getting-started/examples.md

@@ -240,6 +240,29 @@ with Diagram("Broker Consumers", show=False):
     queue = Custom("Message queue", rabbitmq_icon)
 
     queue >> consumers >> Aurora("Database")
-````
+```
+
+![rabbitmq consumers diagram](/img/rabbitmq_consumers_diagram.png)
+
+## Diagram with Custom Edges
+
+```python
+from diagrams import Edge
+from diagrams.aws.compute import EC2
+
+# line
+[EC2("node3"), EC2("node4")] - Edge(color='red', label='label1', style='dotted') - EC2("node")
+
+# list of nodes, one directional
+[EC2("node3"), EC2("node4")] >> Edge(color='red', label='label1', style='dotted') >> EC2("node")
+[EC2("node3"), EC2("node4")] << Edge(color='green', label='label2', style='dashed') << EC2("node")
+
+# both directional
+EC2("node") << Edge(color='blue', label='label3', style='bold') >> EC2("node")
+
+# loop
+node = EC2("node")
+node >> Edge(color='pink', label='label4', style='solid') << node
+```
 
-![rabbitmq consumers diagram](/img/rabbitmq_consumers_diagram.png)
+![custom edges diagram](/img/custom_edges_diagram.png)

docs/guides/edge.md

@@ -0,0 +1,33 @@
+---
+id: edge
+title: Edges
+---
+
+Edge is an object representing an edge between two Nodes.
+
+## Basic
+
+Node is an abstract concept that represents a single system component object. 
+
+An edge object contains three attributes: **label**, **color** and **style** which mirror corresponding graphviz edge attributes.
+
+```python
+from diagrams import Edge
+from diagrams.aws.compute import EC2
+
+# line
+[EC2("node3"), EC2("node4")] - Edge(color='red', label='label1', style='dotted') - EC2("node")
+
+# list of nodes, one directional
+[EC2("node3"), EC2("node4")] >> Edge(color='red', label='label1', style='dotted') >> EC2("node")
+[EC2("node3"), EC2("node4")] << Edge(color='green', label='label2', style='dashed') << EC2("node")
+
+# both directional
+EC2("node") << Edge(color='blue', label='label3', style='bold') >> EC2("node")
+
+# loop
+node = EC2("node")
+node >> Edge(color='pink', label='label4', style='solid') << node
+```
+
+![custom edges diagram](/img/custom_edges_diagram.png)