diff --git a/plesna/__init__.py b/plesna/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/plesna/graph_set.py b/plesna/graph_set.py new file mode 100644 index 0000000..71e3f93 --- /dev/null +++ b/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 diff --git a/tests/graphs/__init__.py b/tests/graphs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/graphs/test_graph_set.py b/tests/graphs/test_graph_set.py new file mode 100644 index 0000000..b4c5fda --- /dev/null +++ b/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])}