feat-custom-edge: Add custom edge class for edge operations

Author: Bishalsarang

tests/test_edge.py

@@ -0,0 +1,19 @@
+from visualiser import Edge, Node
+
+"""
+    Test case class for nodes 
+"""
+
+
+class TestEdge:
+    def test_create_new_edge(self):
+        # Create Node
+        A = Node('A', color='red', style='filled')
+        B = Node('B', label='B label', color='red', style='filled')
+
+        # Create Edge
+        edge = Edge(A, B, label='Test Label')
+
+        assert edge.label == 'Test Label'
+        assert edge.source_node.label == 'A'
+        assert edge.source_node == A

visualiser/__init__.py

@@ -1,5 +1,8 @@
 # Version of the recursion-visualiser
 __version__ = "1.0.1"
 
-from  .node import Node
+# Maintain this order to avoid circular import
+# because we are using Node inside Edge for type annotations.
+from .node import Node
+from .edge import Edge
 from visualiser import *

visualiser/edge.py

@@ -0,0 +1,103 @@
+import copy
+
+
+from visualiser import  Node
+class Edge:
+
+    def __init__(self, source_node: Node, destination_node: Node, label: str = '',
+                 **attrs: str) -> None:
+        # TODO: Remove copy after finding a better way to do this.
+        self._source_node = copy.deepcopy(source_node)
+        self._destination_node = copy.deepcopy(destination_node)
+
+        self._label = label
+        self._attrs = attrs
+
+    @property
+    def label(self) -> str:
+        """
+        Get label for edge.
+        :return: str
+        """
+        return self._label
+
+    @label.setter
+    def label(self, _label: str) -> None:
+        """
+        Set label for edge.
+        :param _label: str
+        """
+        self._label = _label
+
+    @property
+    def source_node(self) -> Node:
+        """
+        Get source node
+        :return: Node
+        """
+        return self._source_node
+
+    @source_node.setter
+    def source_node(self, _source_node: Node) -> None:
+        """
+        Set source node.
+        :param _source_node: Node
+        """
+        self._source_node = _source_node
+
+    @property
+    def destination_node(self) -> Node:
+        """
+        Get destination node.
+        :return: Node
+        """
+        return self._destination_node
+
+    @destination_node.setter
+    def destination_node(self, _destination_node: Node) -> None:
+        """
+        Sets destination node.
+        :param _destination_node: Node
+        """
+        self._destination_node = _destination_node
+
+    def get_attribute(self, key: str) -> str:
+        """
+        Get attribute for edge.
+        :param key: str
+        :return: The value of attribute with key
+        """
+        return self._attrs.get(key)
+
+    def set_attribute(self, key: str, value: str) -> None:
+        """
+        Set attribute for edge
+        :param key: str
+        :param value: str
+        """
+        self._attrs[key] = value
+
+    def remove_attribute(self, key: str) -> None:
+        """
+        Remove attribute from edge.
+        :param key: str
+        """
+        del self._attrs[key]
+
+    def get_attributes_string(self) -> str:
+        """
+        Get attributes string enclosed in []
+        :return:
+        """
+        if len(self._label) == 0:
+            return '[' + ', '.join([f'{key}="{value}"' for key, value in self._attrs.items()]) + ']'
+
+        return '[' + f'label="{self._label}", ' + ', '.join(
+            [f'{key}="{value}"' for key, value in self._attrs.items()]) + ']'
+
+    def to_string(self) -> str:
+        """
+        Converts dot string equivalent of the current edge.
+        :return: Str
+        """
+        return f'{self._source_node.name} -> {self._destination_node.name}  {self.get_attributes_string()}'

visualiser/node.py

@@ -1,7 +1,8 @@
+
 class Node:
     def __init__(self, name, label='', **attrs):
         # TODO: Add support for initialization of attributes from attribute dict
-        self._name = name
+        self._name =  name
         self._attrs = attrs
 
         if len(label) == 0:
@@ -12,6 +13,9 @@ def __init__(self, name, label='', **attrs):
     def __repr__(self):
         return f"Node('{self.name}')"
 
+    def __eq__(self, other):
+        return self.to_string() == other.to_string()
+
     @property
     def name(self):
         return self._name
@@ -29,7 +33,7 @@ def label(self, _label):
         self._label = _label
 
     def get_attribute(self, key):
-        return self._attrs[key]
+        return self._attrs.get(key)
 
     def set_attribute(self, key, value):
         self._attrs[key] = value
@@ -42,4 +46,4 @@ def get_attributes_string(self):
             [f'{key}="{value}"' for key, value in self._attrs.items()]) + '];'
 
     def to_string(self):
-        return f'{self.name} {self.get_attributes_string()}'
+        return f'{self.name} {self.get_attributes_string()}'