Feat: init graphs

plesna/graph_set.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 EdgeOnSet(BaseModel):
+    arrow: Callable
+    sources: dict[str, Node]
+    targets: dict[str, Node]
+    edge_kwrds: dict = {}
+
+
+class GraphSet:
+    def __init__(self):
+        self._edges = []
+        self._node_sets = set()
+
+    def append(self, edge: EdgeOnSet):
+        self._edges.append(edge)
+        self._node_sets.add(frozenset(edge.sources.values()))
+        self._node_sets.add(frozenset(edge.targets.values()))
+
+    @property
+    def node_sets(self):
+        return self._node_sets

tests/graphs/test_graph_set.py

@@ -0,0 +1,18 @@
+from plesna.graph_set import EdgeOnSet, GraphSet, Node
+
+
+def test_init():
+    nodeA = Node(name="A")
+    nodeB = Node(name="B")
+    nodeC = Node(name="C")
+
+    def arrow(sources, targets):
+        targets["C"].infos["res"] = sources["A"].name + sources["B"].name
+
+    edge1 = EdgeOnSet(
+        arrow=arrow, sources={"A": nodeA, "B": nodeB}, targets={"C": nodeC}
+    )
+    graph_set = GraphSet()
+    graph_set.append(edge1)
+
+    assert graph_set.node_sets == {frozenset([nodeA, nodeB]), frozenset([nodeC])}