refact: replace callback with str for arrow in graph_set

This commit is contained in:
Bertrand Benjamin 2025-01-05 06:51:14 +01:00
parent f0315d09b9
commit d1c1b7420d
3 changed files with 8 additions and 14 deletions

View File

@ -1,5 +1,4 @@
from functools import reduce
from typing import Callable
from pydantic import BaseModel

View File

@ -12,10 +12,9 @@ class Node(BaseModel):
class EdgeOnSet(BaseModel):
arrow: Callable
sources: dict[str, Node]
targets: dict[str, Node]
edge_kwrds: dict = {}
arrow: str
sources: list[Node]
targets: list[Node]
class GraphSet:
@ -25,8 +24,8 @@ class GraphSet:
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()))
self._node_sets.add(frozenset(edge.sources))
self._node_sets.add(frozenset(edge.targets))
@property
def node_sets(self):

View File

@ -2,17 +2,13 @@ from plesna.graph.graph_set import EdgeOnSet, GraphSet, Node
def test_init():
graph_set = GraphSet()
nodeA = Node(name="A")
nodeB = Node(name="B")
nodeC = Node(name="C")
edge1 = EdgeOnSet(arrow="arrow", sources=[nodeA, nodeB], targets=[nodeC])
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])}