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

Author: Bishalsarang

tests/test_graph.py

@@ -0,0 +1,35 @@
+from visualiser import Node, Edge, Graph
+
+
+class TestGraph:
+    def test_normal_graph_flow(self):
+        graph = Graph('hello')
+
+        # Create Node
+        A = Node('A', color='red')
+        B = Node('B', color='green')
+        C = Node('C', color='yellow')
+
+        # Add node
+        graph.add_node(A)
+        graph.add_node(B)
+        graph.add_node(C)
+
+        # Make edge
+        edge1 = Edge(A, B)
+        edge2 = Edge(A, C)
+
+        # Add node and graph to edge
+        graph.add_edge(edge1)
+        graph.add_edge(edge2)
+
+        assert graph.to_string() == 'digraph G {\nA [label="A", color="red"];\nB [label="B", color="green"];\nC [label="C", color="yellow"];\nA -> B  [];\nA -> C  [];\n}'
+
+    def test_node_methods(self):
+        pass
+
+    def test_edge_methods(self):
+        pass
+
+    def test_mutations(self):
+        pass

tests/test_node.py

@@ -3,6 +3,8 @@
 """
     Test case class for nodes 
 """
+
+
 class TestNode:
     def test_create_new_node(self):
         node = Node('bishal', 'Bishal label', color='red', style='filled')
@@ -36,5 +38,3 @@ def test_rename_name_label(self):
         assert node.name == 'Bishal renamed'
 
         assert node.to_string() == 'Bishal renamed [label="Bishal renamed label", color="green", style="filled"];'
-
-

visualiser/__init__.py

@@ -5,4 +5,5 @@
 # because we are using Node inside Edge for type annotations.
 from .node import Node
 from .edge import Edge
+from .graph import Graph
 from visualiser import *

visualiser/edge.py

@@ -1,15 +1,18 @@
 import copy
+from typing import Union
+
+from visualiser import Node
 
 
-from visualiser import  Node
 class Edge:
 
-    def __init__(self, source_node: Node, destination_node: Node, label: str = '',
+    def __init__(self, source_node: Union[Node, str], destination_node: Union[Node, str], 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._source_node = Node(source_node) if isinstance(source_node, str) else copy.deepcopy(source_node)
+        self._destination_node = Node(destination_node) if isinstance(destination_node, str) else copy.deepcopy(
+            destination_node)
 
+        self._name = f"{self._source_node.name} -> {self._destination_node.name}"
         self._label = label
         self._attrs = attrs
 
@@ -29,6 +32,22 @@ def label(self, _label: str) -> None:
         """
         self._label = _label
 
+    @property
+    def name(self) -> str:
+        """
+        Get name for edge.
+        :return: str
+        """
+        return self._name
+
+    @name.setter
+    def name(self, _name: str) -> None:
+        """
+        Set label for edge.
+        :param _name: str
+        """
+        self._name = _name
+
     @property
     def source_node(self) -> Node:
         """
@@ -82,18 +101,19 @@ def remove_attribute(self, key: str) -> None:
         Remove attribute from edge.
         :param key: str
         """
-        del self._attrs[key]
+        if self._attrs.get(key):
+            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 '[' + ', '.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()]) + ']'
+            [f'{key}="{value}"' for key, value in self._attrs.items()]) + '];'
 
     def to_string(self) -> str:
         """

visualiser/graph.py

@@ -0,0 +1,102 @@
+import copy
+
+from visualiser import Node, Edge
+
+
+class Graph:
+    def __init__(self, name='', **attrs):
+        self._name = name
+        self._attrs = attrs
+
+        self._nodes = []
+        self._edges = []
+
+    def get_node(self, name):
+        filtered_nodes = list(filter(lambda node: node.name == name, self._nodes))
+
+        return filtered_nodes[0] if len(filtered_nodes) else None
+
+    def remove_node(self, _node):
+        if isinstance(_node, Node):
+            self._nodes = list(filter(lambda node: node.name != _node.name, self._nodes))
+            return
+
+        self._nodes = list(filter(lambda node: node.name != _node, self._nodes))
+
+    def add_node(self, node):
+        if self.get_node(node.name) is not None:
+            self.remove_node(node)
+
+        self._nodes.append(copy.deepcopy(node))
+
+    def set_node_attributes(self, node, **attrs):
+        for key, value in attrs.items():
+            self.set_node_attribute(node, key, value)
+
+    def set_node_attribute(self, _node, key, value):
+        def _set_attribute(name):
+            node = self.get_node(name)
+            if node:
+                node.set_attribute(key, value)
+
+        if isinstance(_node, Node):
+            _set_attribute(_node.name)
+            return
+
+        _set_attribute(_node)
+
+    def remove_node_attribute(self, _node, key):
+        def _remove_attribute(name):
+            node = self.get_node(name)
+            if node:
+                node.remove_attribute(key)
+
+        if isinstance(_node, Node):
+            _remove_attribute(_node.name)
+            return
+
+        _remove_attribute(_node)
+
+    def add_edge(self, edge):
+        if self.get_edge(edge.name) is not None:
+            self.remove_edge(edge)
+
+        # TODO: Check if node used in node exist in graph. If not create one
+        self._edges.append(copy.deepcopy(edge))
+
+    def get_edge(self, name):
+        filtered_edges =  list(filter(lambda edge: edge.name == name, self._edges))
+
+        return filtered_edges[0] if len(filtered_edges) else None
+
+    def remove_edge(self, _edge):
+        if isinstance(_edge, Edge):
+            self._edges = list(filter(lambda edge: edge.name != _edge.name, self._edges))
+            return
+
+        self._edges = list(filter(lambda edge: edge.name != _edge, self._edges))
+
+    def set_edge_label(self, name, value):
+        edge = self.get_edge(name)
+        if edge:
+            edge.label = value
+
+    def highlight_node(self, name, color):
+        self.set_node_attribute(name, 'color', color)
+
+    def highlight_edge(self, name, color):
+        edge = self.get_edge(name)
+        if edge:
+            edge.set_attribute('color', color)
+
+    def reverse_edge_orientation(self, name):
+        pass
+
+    def get_nodes_string(self):
+        return "\n".join(list(map(lambda node: node.to_string(), self._nodes)))
+
+    def get_edges_string(self):
+        return "\n".join(list(map(lambda edge: edge.to_string(), self._edges)))
+
+    def to_string(self):
+        return "digraph G {\n" + f"{self.get_nodes_string()}\n" + f"{self.get_edges_string()}\n" + "}"

visualiser/node.py

@@ -42,6 +42,7 @@ def remove_attribute(self, key):
         del self._attrs[key]
 
     def get_attributes_string(self):
+        # TODO: Check if no any attributes and labels to avoid '[]'
         return '[' + f'label="{self._label}", ' + ', '.join(
             [f'{key}="{value}"' for key, value in self._attrs.items()]) + '];'