From c90f407cfc2a09c78d9be6d44070982df3288be2 Mon Sep 17 00:00:00 2001
From: Bertrand Benjamin <benjamin.bertrand@opytex.org>
Date: Mon, 7 Oct 2024 06:09:01 +0200
Subject: [PATCH] Feat: init graphs

---
 plesna/__init__.py             |  0
 plesna/graph_set.py            | 33 +++++++++++++++++++++++++++++++++
 tests/graphs/__init__.py       |  0
 tests/graphs/test_graph_set.py | 18 ++++++++++++++++++
 4 files changed, 51 insertions(+)
 create mode 100644 plesna/__init__.py
 create mode 100644 plesna/graph_set.py
 create mode 100644 tests/graphs/__init__.py
 create mode 100644 tests/graphs/test_graph_set.py

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])}