feat: update interface

Author: brian14708

pyproject.toml

@@ -61,3 +61,4 @@ venvPath = ".venv"
 reportAny = false
 reportExplicitAny = false
 reportUnreachable = false
+reportImportCycles = false

src/duron/__init__.py

@@ -1,5 +1,4 @@
-from duron.context import get_context
+from duron.context import Context, get_context
 from duron.fn import durable
-from duron.task import task
 
-__all__ = ["durable", "task", "get_context"]
+__all__ = ["durable", "get_context", "Context"]

src/duron/fn.py

@@ -3,18 +3,26 @@
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
+    Any,
+    Concatenate,
     Generic,
     ParamSpec,
     TypeVar,
+    cast,
 )
 
 from duron.codec import DefaultCodec
+from duron.task import Task, TaskGuard
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Callable, Coroutine
 
     from duron.codec import Codec
+    from duron.context import Context
+    from duron.log import LogStorage
 
+    _TOffset = TypeVar("_TOffset")
+    _TLease = TypeVar("_TLease")
 
 _T_co = TypeVar("_T_co", covariant=True)
 _P = ParamSpec("_P")
@@ -23,21 +31,29 @@
 @dataclass(slots=True)
 class DurableFn(Generic[_P, _T_co]):
     codec: Codec
-    fn: Callable[_P, _T_co]
+    fn: Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]]
 
-    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _T_co:
-        return self.fn(*args, **kwargs)
+    def __call__(
+        self,
+        log: LogStorage[_TOffset, _TLease],
+    ) -> TaskGuard[_P, _T_co]:
+        return TaskGuard(Task(self, cast("LogStorage[object, object]", log)))
 
 
 def durable(
     *, codec: Codec | None = None
-) -> Callable[[Callable[_P, _T_co]], DurableFn[_P, _T_co]]:
+) -> Callable[
+    [Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]]],
+    DurableFn[_P, _T_co],
+]:
     """
     Mark a function as durable, meaning its execution can be recorded and
     replayed.
     """
 
-    def decorate(fn: Callable[_P, _T_co]) -> DurableFn[_P, _T_co]:
+    def decorate(
+        fn: Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]],
+    ) -> DurableFn[_P, _T_co]:
         return DurableFn(codec=codec or DefaultCodec(), fn=fn)
 
     return decorate

src/duron/task.py

@@ -38,22 +38,13 @@
     )
     from duron.ops import Op
 
-    _TOffset = TypeVar("_TOffset")
-    _TLease = TypeVar("_TLease")
 
 _T = TypeVar("_T")
 _P = ParamSpec("_P")
 
 _CURRENT_VERSION = 0
 
 
-def task(
-    task_co: DurableFn[_P, Coroutine[Any, Any, _T]],
-    log: LogStorage[_TOffset, _TLease],
-) -> TaskGuard[_P, _T]:
-    return TaskGuard(Task(task_co, log))
-
-
 @final
 class TaskGuard(Generic[_P, _T]):
     def __init__(self, task: Task[_P, _T]) -> None:
@@ -74,9 +65,8 @@ async def __aexit__(
 class Task(Generic[_P, _T]):
     def __init__(
         self,
-        task_fn: DurableFn[_P, Coroutine[Any, Any, _T]],
-        log: LogStorage[_TOffset, _TLease],
-        codec: Codec | None = None,
+        task_fn: DurableFn[_P, _T],
+        log: LogStorage[object, object],
     ) -> None:
         self._task_fn = task_fn
         self._log = log
@@ -91,7 +81,7 @@ async def start(self, *args: _P.args, **kwargs: _P.kwargs) -> None:
         }
         self._run = _TaskRun(
             _task_prelude(self._task_fn, lambda: init),
-            cast("LogStorage[object, object]", self._log),
+            self._log,
             codec,
         )
         await self._run.resume()
@@ -103,7 +93,7 @@ def cb() -> TaskInitParams:
         task = _task_prelude(self._task_fn, cb)
         self._run = _TaskRun(
             task,
-            cast("LogStorage[object, object]", self._log),
+            self._log,
             self._task_fn.codec,
         )
         await self._run.resume()
@@ -121,7 +111,7 @@ class TaskInitParams(TypedDict):
 
 
 async def _task_prelude(
-    task_fn: DurableFn[..., Coroutine[Any, Any, object]],
+    task_fn: DurableFn[..., object],
     init: Callable[[], TaskInitParams],
 ) -> object:
     ctx = get_context()
@@ -131,7 +121,7 @@ async def _task_prelude(
     codec = task_fn.codec
     args = (codec.decode_json(arg) for arg in init_params["args"])
     kwargs = {k: codec.decode_json(v) for k, v in init_params["kwargs"].items()}
-    return await task_fn.fn(*args, **kwargs)
+    return await task_fn.fn(get_context(), *args, **kwargs)
 
 
 @final

tests/test_task.py

@@ -6,12 +6,16 @@
 import random
 import uuid
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import pytest
 
-from duron import durable, get_context, task
+from duron import durable
 from duron.log.storage import MemoryLogStorage
 
+if TYPE_CHECKING:
+    from duron.context import Context
+
 
 @pytest.mark.asyncio
 async def test_task():
@@ -21,8 +25,7 @@ async def u() -> str:
         return str(uuid.uuid4())
 
     @durable()
-    async def activity(i: str) -> str:
-        ctx = get_context()
+    async def activity(ctx: Context, i: str) -> str:
         x = await asyncio.gather(
             ctx.run(u),
             ctx.run(u),
@@ -44,18 +47,18 @@ async def activity(i: str) -> str:
     }
 
     log = MemoryLogStorage()
-    async with task(activity, log) as t:
+    async with activity(log) as t:
         await t.start("test")
         a = await t.wait()
     assert set(e["id"] for e in await log.entries()) == IDS
 
-    async with task(activity, log) as t:
+    async with activity(log) as t:
         await t.start("test")
         b = await t.wait()
     assert a == b
 
     log2 = MemoryLogStorage((await log.entries())[:-2])
-    async with task(activity, log2) as t:
+    async with activity(log2) as t:
         await t.start("test")
         c = await t.wait()
     assert a == c
@@ -65,8 +68,7 @@ async def activity(i: str) -> str:
 @pytest.mark.asyncio
 async def test_task_error():
     @durable()
-    async def activity():
-        ctx = get_context()
+    async def activity(ctx: Context):
         _ = await ctx.run(lambda: asyncio.sleep(0.1))
 
         async def error():
@@ -76,11 +78,11 @@ async def error():
 
     log = MemoryLogStorage()
     with pytest.raises(check=lambda v: "test error" in str(v)):
-        async with task(activity, log) as t:
+        async with activity(log) as t:
             await t.start()
             await t.wait()
     with pytest.raises(check=lambda v: "test error" in str(v)):
-        async with task(activity, log) as t:
+        async with activity(log) as t:
             await t.start()
             await t.wait()
 
@@ -90,18 +92,17 @@ async def test_resume():
     sleep = 9999
 
     @durable()
-    async def activity(s: str) -> str:
-        ctx = get_context()
+    async def activity(ctx: Context, s: str) -> str:
         _ = await ctx.run(lambda: asyncio.sleep(sleep))
         return s
 
     log = MemoryLogStorage()
-    async with task(activity, log) as t:
+    async with activity(log) as t:
         await t.start("hello")
         with pytest.raises(asyncio.TimeoutError):
             _ = await asyncio.wait_for(t.wait(), 0.1)
 
-    async with task(activity, log) as t:
+    async with activity(log) as t:
         sleep = 0
         await t.resume()
         x = await t.wait()
@@ -127,13 +128,12 @@ def decode_json(self, encoded: object) -> object:
 @pytest.mark.asyncio
 async def test_serialize():
     @durable(codec=PickleCodec())
-    async def activity() -> CustomPoint:
-        ctx = get_context()
+    async def activity(ctx: Context) -> CustomPoint:
         pt = await ctx.run(lambda: CustomPoint(x=1, y=2))
         return CustomPoint(x=pt.x + 5, y=pt.y + 10)
 
     log = MemoryLogStorage()
-    async with task(activity, log) as t:
+    async with activity(log) as t:
         await t.start()
         a = await t.wait()
         assert type(a) is CustomPoint