tinyfan 0.2.2__tar.gz → 0.2.3__tar.gz

{tinyfan-0.2.2 → tinyfan-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinyfan
-Version: 0.2.2
+Version: 0.2.3
 Summary: Tiny Argo Workflows python devkit
 Project-URL: Documentation, https://github.com/eunchuldev/tinyfan#readme
 Project-URL: Issues, https://github.com/eunchuldev/tinyfan/issues

{tinyfan-0.2.2 → tinyfan-0.2.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tinyfan"
-version = "0.2.2"
+version = "0.2.3"
 description = 'Tiny Argo Workflows python devkit'
 readme = "README.md"
 requires-python = ">=3.12"

{tinyfan-0.2.2 → tinyfan-0.2.3}/src/tinyfan/codegen.py

@@ -191,13 +191,13 @@ class AssetTree:
 
         relatives_by_schedules = self.relatives_by_schedules()
         g: nx.DiGraph = nx.DiGraph()
-        # g.add_node(self.flow.name)
-        # g.add_nodes_from([k for k in relatives_by_schedules.keys()])
-        # g.add_nodes_from(self.nodes.keys())
-        for (_, schedule), relatives in relatives_by_schedules.items():
-            g.add_edge(self.flow.name, schedule)
-            g.add_edges_from([(schedule, r.asset.name) for r in relatives])
-            g.add_edges_from([(r.asset.name, d.asset.name) for r in relatives for d in r.parents])
+        g.add_node(self.flow.name, bullet="■")
+        for (tz, schedule), relatives in relatives_by_schedules.items():
+            schedule_node_name = schedule + f" ({tz})"
+            g.add_node(schedule_node_name, bullet="○")
+            g.add_edge(self.flow.name, schedule_node_name)
+            g.add_edges_from([(schedule_node_name, r.asset.name) for r in relatives])
+            g.add_edges_from([(d.asset.name, r.asset.name) for r in relatives for d in r.parents])
         return dagviz(g, round_angle=True)
 
     def compile(

{tinyfan-0.2.2 → tinyfan-0.2.3}/src/tinyfan/dagviz.py

@@ -15,6 +15,7 @@ def dagviz(g: nx.DiGraph, round_angle: bool = False) -> str:
     assert nx.is_directed_acyclic_graph(g), "Graph contains cycles"
 
     rows: List[Hashable] = []
+    bullets: List[str] = []
     node_to_row: Dict[Hashable, int] = {}
     indents: List[int] = []
 
@@ -27,6 +28,7 @@ def dagviz(g: nx.DiGraph, round_angle: bool = False) -> str:
         for i in range(len(sources)):
             node_to_row[sources[i]] = len(rows)
             rows.append(sources[i])
+            bullets.append(g.nodes[sources[i]].get("bullet", "•"))
             indents.append(indent + i)
         _process_dag(nx.DiGraph(g.subgraph(set(g.nodes).difference(sources))), indent=indent + len(sources))
 
@@ -50,5 +52,5 @@ def dagviz(g: nx.DiGraph, round_angle: bool = False) -> str:
         for k in range(indents[i] * 2 + 1, indents[n] * 2):
             a[n][k] = "─"
 
-    lines: List[str] = [x.tounicode() + "• " + str(i).replace("\n", " ") for x, i in zip(a, rows)]
+    lines: List[str] = [x.tounicode() + f"{b} " + str(i).replace("\n", " ") for x, b, i in zip(a, bullets, rows)]
     return "\n".join(lines)

{tinyfan-0.2.2 → tinyfan-0.2.3}/src/tinyfan/flow.py

@@ -1,9 +1,10 @@
 import inspect
+import asyncio
 from .stores.base import StoreBase
 from .stores.naive import NaiveStore
 from .flowrundata import FlowRunData, StoreIdx, UMeta
 from dataclasses import dataclass, field
-from typing import Callable, Any, Generic, TypeVar
+from typing import Callable, Any, Generic, TypeVar, Awaitable, Coroutine, cast
 from .resources.base import ResourceBase
 from .argo_typing import ScriptTemplate
 from .config import Config, ConfigValue
@@ -41,7 +42,7 @@ DEFAULT_FLOW: Flow = Flow(
 @dataclass
 class Asset(Generic[Ret, UMeta, StoreIdx]):
     flow: Flow
-    func: Callable[..., Ret]
+    func: Callable[..., Ret | Awaitable[Ret]]
     store: StoreBase[Ret, UMeta, StoreIdx]
     schedule: str | None = None
     tz: str | None = None
@@ -84,11 +85,15 @@ class Asset(Generic[Ret, UMeta, StoreIdx]):
             if name in func_param_names:
                 params[name] = data
         ret = self.func(**params)
-        store_entry_idx = self.store.store(ret, rundata)
+        if inspect.isawaitable(ret):
+            res: Ret = asyncio.run(cast(Coroutine[Any, Any, Ret], ret))
+        else:
+            res = cast(Ret, ret)
+        store_entry_idx = self.store.store(res, rundata)
         rundata["store_entry_idx"] = store_entry_idx
         if "parents" in rundata:
             del rundata["parents"]
-        return (ret, rundata)
+        return (res, rundata)
 
 
 class AssetFunc(Generic[Ret, UMeta, StoreIdx]):

tinyfan-0.2.3/tests/test_decorators.py

@@ -0,0 +1,32 @@
+from tinyfan import asset
+
+
+def test_preserve_original_functionality():
+    @asset()
+    def test(arg: str):
+        return arg
+
+    assert test(arg="test") == "test"
+
+
+def test_asset_run():
+    @asset()
+    def test(ds: str):
+        return ds
+
+    actual, _ = test.asset.run({"ds": "2020-01-01"})
+    expected = "2020-01-01"
+    assert actual == expected
+
+
+def test_async_asset_run():
+    async def awaitable() -> str:
+        return " awaited"
+
+    @asset()
+    async def test(ds: str):
+        return ds + await awaitable()
+
+    actual, _ = test.asset.run({"ds": "2020-01-01"})
+    expected = "2020-01-01 awaited"
+    assert actual == expected

{tinyfan-0.2.2 → tinyfan-0.2.3}/uv.lock

@@ -784,7 +784,7 @@ wheels = [
 
 [[package]]
 name = "tinyfan"
-version = "0.2.2"
+version = "0.2.3"
 source = { editable = "." }
 dependencies = [
     { name = "croniter" },

tinyfan-0.2.2/tests/test_decorators.py

@@ -1,9 +0,0 @@
-from tinyfan import asset
-
-
-def test_preserve_original_functionality():
-    @asset()
-    def test(arg: str):
-        return arg
-
-    assert test(arg="test") == "test"