diff --git a/plesna/graph.py b/plesna/graph.py new file mode 100644 index 0000000..00b8aff --- /dev/null +++ b/plesna/graph.py @@ -0,0 +1,33 @@ +from typing import Callable + +from pydantic import BaseModel + + +class Node(BaseModel): + name: str + infos: dict = {} + + def __hash__(self): + return hash(self.name) + + +class Edge(BaseModel): + arrow_name: str + source: Node + target: Node + edge_kwrds: dict = {} + + +class Graph: + def __init__(self): + self._edges = [] + self._nodes = set() + + def append(self, edge: Edge): + self._edges.append(edge) + self._nodes.add(edge.source) + self._nodes.add(edge.target) + + @property + def nodes(self): + return self._nodes diff --git a/tests/graphs/test_graph.py b/tests/graphs/test_graph.py new file mode 100644 index 0000000..1fbe18b --- /dev/null +++ b/tests/graphs/test_graph.py @@ -0,0 +1,16 @@ +from plesna.graph import Edge, Graph, Node + + +def test_init(): + nodeA = Node(name="A") + nodeB = Node(name="B") + nodeC = Node(name="C") + + edge1 = Edge(arrow_name="arrow", source=nodeA, target=nodeC) + edge2 = Edge(arrow_name="arrow", source=nodeB, target=nodeC) + + graph = Graph() + graph.append(edge1) + graph.append(edge2) + + assert graph.nodes == {nodeA, nodeB, nodeC}