From 0ae64392176e681ea7d2c26ff1dce49494bc1f77 Mon Sep 17 00:00:00 2001
From: Bertrand Benjamin <benjamin.bertrand@opytex.org>
Date: Sun, 5 Jan 2025 18:40:19 +0100
Subject: [PATCH] refact: replace edge_kwrds with metadata in models

---
 plesna/graph/graph_set.py                |  2 +-
 plesna/models/graphs.py                  |  5 +++--
 plesna/models/libs/flux_graph.py         | 22 ++++++++++++++++++++++
 tests/dataplatform/test_dataplateform.py | 10 +++++-----
 4 files changed, 31 insertions(+), 8 deletions(-)
 create mode 100644 plesna/models/libs/flux_graph.py

diff --git a/plesna/graph/graph_set.py b/plesna/graph/graph_set.py
index c9e71b9..58993a4 100644
--- a/plesna/graph/graph_set.py
+++ b/plesna/graph/graph_set.py
@@ -28,7 +28,7 @@ class GraphSet:
             graph.add_nodes(node_set)
         for edge in self._edges:
             flatten_edge = [
-                Edge(arrow=edge.arrow, source=s, target=t, edge_kwrds=edge.edge_kwrds)
+                Edge(arrow=edge.arrow, source=s, target=t, metadata=edge.metadata)
                 for (s, t) in product(edge.sources, edge.targets)
             ]
             graph.add_edges(flatten_edge)
diff --git a/plesna/models/graphs.py b/plesna/models/graphs.py
index 3506005..f0812e2 100644
--- a/plesna/models/graphs.py
+++ b/plesna/models/graphs.py
@@ -3,6 +3,7 @@ from pydantic import BaseModel
 
 class Node(BaseModel):
     name: str
+    metadata: dict = {}
 
     def __hash__(self):
         return hash(self.name)
@@ -12,11 +13,11 @@ class Edge(BaseModel):
     arrow: str
     source: Node
     target: Node
-    edge_kwrds: dict = {}
+    metadata: dict = {}
 
 
 class EdgeOnSet(BaseModel):
     arrow: str
     sources: list[Node]
     targets: list[Node]
-    edge_kwrds: dict = {}
+    metadata: dict = {}
diff --git a/plesna/models/libs/flux_graph.py b/plesna/models/libs/flux_graph.py
new file mode 100644
index 0000000..e070104
--- /dev/null
+++ b/plesna/models/libs/flux_graph.py
@@ -0,0 +1,22 @@
+from collections.abc import Callable
+from plesna.models.flux import Flux
+from plesna.models.graphs import EdgeOnSet, Node
+
+
+def flux_to_edgeonset(
+    flux: Flux,
+    meta_arrow: Callable = lambda flux: flux.id,
+    meta_table: Callable = lambda table: Node(name=table.id),
+) -> EdgeOnSet:
+    """Convert a flux to an EdgeOnSet
+
+    :param flux: the flux
+    :meta_arrow: function on flux which returns things to store in arrow field
+    :meta_table: function on table which returns a Node
+
+    """
+    return EdgeOnSet(
+        arrow=meta_arrow(flux),
+        sources=[meta_table(s) for s in flux.sources],
+        targets=[meta_table(t) for t in flux.targets],
+    )
diff --git a/tests/dataplatform/test_dataplateform.py b/tests/dataplatform/test_dataplateform.py
index 64ee7e5..da8a4d5 100644
--- a/tests/dataplatform/test_dataplateform.py
+++ b/tests/dataplatform/test_dataplateform.py
@@ -156,13 +156,13 @@ def test_content_from_graphset(dataplatform: DataPlateform):
             arrow="foo_flux",
             sources=[Node(name="test-raw-username"), Node(name="test-raw-recovery")],
             targets=[Node(name="test-bronze-foo")],
-            edge_kwrds={},
+            metadata={},
         ),
         EdgeOnSet(
             arrow="copy_flux",
             sources=[Node(name="test-raw-username")],
             targets=[Node(name="test-bronze-username")],
-            edge_kwrds={},
+            metadata={},
         ),
     ]
 
@@ -181,19 +181,19 @@ def test_content_from_graph(dataplatform: DataPlateform):
             arrow="foo_flux",
             source=Node(name="test-raw-username"),
             target=Node(name="test-bronze-foo"),
-            edge_kwrds={},
+            metadata={},
         ),
         Edge(
             arrow="foo_flux",
             source=Node(name="test-raw-recovery"),
             target=Node(name="test-bronze-foo"),
-            edge_kwrds={},
+            metadata={},
         ),
         Edge(
             arrow="copy_flux",
             source=Node(name="test-raw-username"),
             target=Node(name="test-bronze-username"),
-            edge_kwrds={},
+            metadata={},
         ),
     ]