refactor: move modules

Author: brian14708

src/duron/__init__.py

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

src/duron/log/codec.py src/duron/codec.pysrc/duron/log/codec.py renamed to src/duron/codec.py

@@ -16,7 +16,7 @@
     pass
 
 if TYPE_CHECKING:
-    from duron.log.entry import JSONValue
+    from duron.log import JSONValue
 
 __all__ = ["Codec", "DEFAULT_CODEC"]
 

src/duron/mark.py src/duron/fn.pysrc/duron/mark.py renamed to src/duron/fn.py

@@ -9,38 +9,40 @@
     cast,
 )
 
-from duron.log.codec import DEFAULT_CODEC
+from duron.codec import DEFAULT_CODEC
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from duron.log.codec import Codec
+    from duron.codec import Codec
 
 
 _T_co = TypeVar("_T_co", covariant=True)
 _P = ParamSpec("_P")
 
 
-@dataclass
-class DurableFnParams:
-    codec: Codec
+class Fn(Protocol[_P, _T_co]):
+    __duron__: FnOptions
 
+    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _T_co: ...
 
-class DurableFn(Protocol[_P, _T_co]):
-    __durable__: DurableFnParams
 
-    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _T_co: ...
+@dataclass(slots=True)
+class FnOptions:
+    codec: Codec
 
 
-def durable() -> Callable[[Callable[_P, _T_co]], DurableFn[_P, _T_co]]:
+def fn(
+    *, codec: Codec = DEFAULT_CODEC
+) -> Callable[[Callable[_P, _T_co]], Fn[_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]:
-        d = cast("DurableFn[_P, _T_co]", fn)
-        d.__durable__ = DurableFnParams(codec=DEFAULT_CODEC)
+    def decorate(fn: Callable[_P, _T_co]) -> Fn[_P, _T_co]:
+        d = cast("Fn[_P, _T_co]", fn)
+        d.__duron__ = FnOptions(codec=codec)
         return d
 
     return decorate

src/duron/log/__init__.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Generic, Protocol
+
+from typing_extensions import TypedDict, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+    from typing import Literal, TypeGuard
+
+    from typing_extensions import NotRequired
+
+
+__all__ = ["LogStorage", "Entry", "AnyEntry", "is_entry"]
+
+_TOffset = TypeVar("_TOffset")
+_TLease = TypeVar("_TLease")
+JSONValue = dict[str, "JSONValue"] | list["JSONValue"] | str | int | float | bool | None
+
+
+class _BaseEntry(TypedDict):
+    id: str
+    # Unix timestamp in microseconds
+    ts: int
+    meta: NotRequired[dict[str, str]]
+
+
+class ErrorInfo(TypedDict):
+    code: int
+    message: str
+    state: NotRequired[str]  # opaque
+
+
+class PromiseCreateEntry(_BaseEntry):
+    type: Literal["promise/create"]
+
+
+class PromiseCompleteEntry(_BaseEntry):
+    type: Literal["promise/complete"]
+    promise_id: str
+    result: NotRequired[JSONValue]
+    error: NotRequired[ErrorInfo]
+
+
+class StreamCreateEntry(_BaseEntry):
+    type: Literal["stream/create"]
+
+
+class StreamEmitEntry(_BaseEntry):
+    type: Literal["stream/emit"]
+    stream_id: str
+    value: NotRequired[JSONValue]
+    state: NotRequired[str]  # opaque
+
+
+class StreamCompleteEntry(_BaseEntry):
+    type: Literal["stream/complete"]
+    stream_id: str
+    error: NotRequired[ErrorInfo]
+
+
+class AnyEntry(_BaseEntry):
+    type: str
+
+
+Entry = (
+    PromiseCreateEntry
+    | PromiseCompleteEntry
+    | StreamCreateEntry
+    | StreamEmitEntry
+    | StreamCompleteEntry
+)
+
+
+def is_entry(entry: Entry | AnyEntry) -> TypeGuard[Entry]:
+    return entry["type"] in {
+        "promise/create",
+        "promise/complete",
+        "stream/create",
+        "stream/emit",
+        "stream/complete",
+    }
+
+
+class LogStorage(Protocol, Generic[_TOffset, _TLease]):
+    def stream(
+        self, start: _TOffset | None, live: bool
+    ) -> AsyncGenerator[tuple[_TOffset, AnyEntry], None]: ...
+
+    async def acquire_lease(self) -> _TLease: ...
+
+    async def release_lease(self, token: _TLease): ...
+
+    async def append(self, token: _TLease, entry: Entry): ...
+
+    async def flush(self, token: _TLease): ...

src/duron/log/entry.py

@@ -8,14 +8,12 @@
 
 from typing_extensions import override
 
-from .base import LogStorage
+from . import LogStorage
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
 
-    from duron.log.entry import AnyEntry
-
-    from ..entry import Entry
+    from . import AnyEntry, Entry
 
 
 class FileLogStorage(LogStorage[int, str]):

src/duron/log/storage/simple.py src/duron/log/storage.pysrc/duron/log/storage/simple.py renamed to src/duron/log/storage.py

@@ -19,23 +19,23 @@
 
 from duron.context import get_context
 from duron.event_loop import create_loop
-from duron.log.entry import is_entry
+from duron.log import is_entry
 from duron.ops import FnCall, TaskRun
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Coroutine
     from types import TracebackType
 
+    from duron.codec import Codec
     from duron.event_loop import WaitSet
-    from duron.log.codec import Codec
-    from duron.log.entry import (
+    from duron.fn import Fn
+    from duron.log import (
         Entry,
         ErrorInfo,
         JSONValue,
+        LogStorage,
         PromiseCompleteEntry,
     )
-    from duron.log.storage import LogStorage
-    from duron.mark import DurableFn
     from duron.ops import Op
 
     _TOffset = TypeVar("_TOffset")
@@ -48,7 +48,7 @@
 
 
 def task(
-    task_co: DurableFn[_P, Coroutine[Any, Any, _T]],
+    task_co: Fn[_P, Coroutine[Any, Any, _T]],
     log: LogStorage[_TOffset, _TLease],
 ) -> TaskGuard[_P, _T]:
     return TaskGuard(Task(task_co, log))
@@ -74,7 +74,7 @@ async def __aexit__(
 class Task(Generic[_P, _T]):
     def __init__(
         self,
-        task_fn: DurableFn[_P, Coroutine[Any, Any, _T]],
+        task_fn: Fn[_P, Coroutine[Any, Any, _T]],
         log: LogStorage[_TOffset, _TLease],
         codec: Codec | None = None,
     ) -> None:
@@ -83,7 +83,7 @@ def __init__(
         self._run: _TaskRun | None = None
 
     async def start(self, *args: _P.args, **kwargs: _P.kwargs) -> None:
-        codec = self._task_fn.__durable__.codec
+        codec = self._task_fn.__duron__.codec
         init: TaskInitParams = {
             "version": _CURRENT_VERSION,
             "args": [codec.encode_json(arg) for arg in args],
@@ -104,7 +104,7 @@ def cb() -> TaskInitParams:
         self._run = _TaskRun(
             task,
             cast("LogStorage[object, object]", self._log),
-            self._task_fn.__durable__.codec,
+            self._task_fn.__duron__.codec,
         )
         await self._run.resume()
 
@@ -121,14 +121,14 @@ class TaskInitParams(TypedDict):
 
 
 async def _task_prelude(
-    task_fn: DurableFn[..., Coroutine[Any, Any, object]],
+    task_fn: Fn[..., Coroutine[Any, Any, object]],
     init: Callable[[], TaskInitParams],
 ) -> object:
     ctx = get_context()
     init_params = await ctx.run(init)
     if init_params["version"] != _CURRENT_VERSION:
         raise Exception("version mismatch")
-    codec = task_fn.__durable__.codec
+    codec = task_fn.__duron__.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(*args, **kwargs)

src/duron/log/storage/__init__.py

@@ -8,11 +8,10 @@
 import pytest
 from typing_extensions import TypeVar
 
-from duron.log.storage.simple import FileLogStorage, MemoryLogStorage
+from duron.log.storage import FileLogStorage, MemoryLogStorage
 
 if TYPE_CHECKING:
-    from duron.log.entry import Entry
-    from duron.log.storage import LogStorage
+    from duron.log import Entry, LogStorage
 
     _TOffset = TypeVar("_TOffset")
     _TLease = TypeVar("_TLease")