From 9d45625a5ed017fba77cf77a37e137cca0b9578f Mon Sep 17 00:00:00 2001
From: Bertrand Benjamin <benjamin.bertrand@opytex.org>
Date: Wed, 6 Nov 2024 07:00:15 +0100
Subject: [PATCH] Feat: move to models and add consume_flux

---
 plesna/compute/__init__.py         |  0
 plesna/compute/consume_flux.py     |  8 +++++++
 plesna/models/__init__.py          |  0
 plesna/models/flux.py              | 14 ++++++++++++
 plesna/models/transformation.py    | 15 +++++++++++++
 tests/compute/__init__.py          |  0
 tests/compute/test_consume_flux.py | 35 ++++++++++++++++++++++++++++++
 7 files changed, 72 insertions(+)
 create mode 100644 plesna/compute/__init__.py
 create mode 100644 plesna/compute/consume_flux.py
 create mode 100644 plesna/models/__init__.py
 create mode 100644 plesna/models/flux.py
 create mode 100644 plesna/models/transformation.py
 create mode 100644 tests/compute/__init__.py
 create mode 100644 tests/compute/test_consume_flux.py

diff --git a/plesna/compute/__init__.py b/plesna/compute/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/plesna/compute/consume_flux.py b/plesna/compute/consume_flux.py
new file mode 100644
index 0000000..2a176ea
--- /dev/null
+++ b/plesna/compute/consume_flux.py
@@ -0,0 +1,8 @@
+from plesna.models.flux import Flux, FluxMetaData
+
+
+def consume_flux(flux: Flux) -> FluxMetaData:
+    metadata = flux.transformation.function(
+        sources=flux.sources, targets=flux.targets, **flux.transformation.extra_kwrds
+    )
+    return FluxMetaData(data=metadata)
diff --git a/plesna/models/__init__.py b/plesna/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/plesna/models/flux.py b/plesna/models/flux.py
new file mode 100644
index 0000000..28dcc8f
--- /dev/null
+++ b/plesna/models/flux.py
@@ -0,0 +1,14 @@
+from pydantic import BaseModel
+
+from plesna.models.storage import Table
+from plesna.models.transformation import Transformation
+
+
+class Flux(BaseModel):
+    sources: dict[str, Table]
+    targets: dict[str, Table]
+    transformation: Transformation
+
+
+class FluxMetaData(BaseModel):
+    data: dict
diff --git a/plesna/models/transformation.py b/plesna/models/transformation.py
new file mode 100644
index 0000000..5d669a3
--- /dev/null
+++ b/plesna/models/transformation.py
@@ -0,0 +1,15 @@
+from collections.abc import Callable
+
+from pydantic import BaseModel
+
+
+class Transformation(BaseModel):
+    """
+    The function have to have at least 2 arguments: sources and targets
+    Other arguments will came throught extra_kwrds
+
+    The function will have to return metadata as dict
+    """
+
+    function: Callable
+    extra_kwrds: dict = {}
diff --git a/tests/compute/__init__.py b/tests/compute/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/compute/test_consume_flux.py b/tests/compute/test_consume_flux.py
new file mode 100644
index 0000000..9255f4d
--- /dev/null
+++ b/tests/compute/test_consume_flux.py
@@ -0,0 +1,35 @@
+from plesna.compute.consume_flux import consume_flux
+from plesna.models.flux import Flux
+from plesna.models.storage import Table
+from plesna.models.transformation import Transformation
+
+
+def test_consume_flux():
+    sources = {
+        "src1": Table(id="src1", value="here"),
+        "src2": Table(id="src2", value="here"),
+    }
+    targets = {
+        "tgt1": Table(id="tgt1", value="this"),
+        "tgt2": Table(id="tgt2", value="that"),
+    }
+
+    def func(sources, targets, **kwrds):
+        return {
+            "sources": len(sources),
+            "targets": len(targets),
+            "kwrds": len(kwrds),
+        }
+
+    flux = Flux(
+        sources=sources,
+        targets=targets,
+        transformation=Transformation(function=func, extra_kwrds={"extra": "super"}),
+    )
+
+    meta = consume_flux(flux)
+    assert meta.data == {
+        "sources": 2,
+        "targets": 2,
+        "kwrds": 1,
+    }