Feat: add pure graph
This commit is contained in:
parent
c90f407cfc
commit
5c69bb5503
33
plesna/graph.py
Normal file
33
plesna/graph.py
Normal file
@ -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
|
16
tests/graphs/test_graph.py
Normal file
16
tests/graphs/test_graph.py
Normal file
@ -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}
|
Loading…
Reference in New Issue
Block a user