feat: simplify types

Author: brian14708

src/duron/_core/invoke.py

@@ -90,7 +90,7 @@ def __init__(
         self._watchers: list[
             tuple[
                 dict[str, str],
-                StreamObserver[object],
+                StreamObserver,
             ]
         ] = []
 
@@ -207,7 +207,7 @@ def watch_stream(
         observer: ObserverStream[_T_co, None] = ObserverStream()
         self._watchers.append((
             labels,
-            cast("StreamObserver[object]", cast("StreamObserver[_T_co]", observer)),
+            cast("ObserverStream[object, None]", observer),
         ))
         return observer
 
@@ -334,7 +334,7 @@ def __init__(
         log: LogStorage,
         codec: Codec,
         *,
-        watchers: list[tuple[dict[str, str], StreamObserver[object]]] | None = None,
+        watchers: list[tuple[dict[str, str], StreamObserver]] | None = None,
     ) -> None:
         self._loop = create_loop(asyncio.get_running_loop())
         self._task = self._loop.create_task(task)
@@ -518,7 +518,7 @@ async def enqueue_log(
             offset = await self._log.append(self._lease, entry)
             await self.handle_message(offset, entry)
 
-    async def enqueue_op(self, id_: str, fut: OpFuture[object]) -> None:
+    async def enqueue_op(self, id_: str, fut: OpFuture) -> None:
         op = cast("Op", fut.params)
         match op:
             case FnCall():
@@ -566,7 +566,7 @@ async def cb() -> None:
                     entry["ts"] = self.now()
                     await self.enqueue_log(entry)
 
-                def done(f: OpFuture[object]) -> None:
+                def done(f: OpFuture) -> None:
                     if f.cancelled():
                         sid = f.id
                         if self._task_manager.has_pending(sid):

src/duron/_core/ops.py

@@ -2,8 +2,14 @@
 
 import asyncio
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
-from typing_extensions import dataclass_transform, overload
+from typing import TYPE_CHECKING, Final, cast
+from typing_extensions import (
+    Any,
+    Protocol,
+    TypeVar,
+    dataclass_transform,
+    overload,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Coroutine, Mapping
@@ -20,9 +26,7 @@ def frozen(_cls: type[_T]) -> type[_T]: ...
 else:
     frozen = dataclass(slots=True)
 
-_In_contra = TypeVar("_In_contra", contravariant=True)
-
-_EMPTY_DICT: dict[str, Any] = {}
+_EMPTY_DICT: Final = cast("dict[str, Any]", {})
 
 
 @frozen
@@ -36,9 +40,7 @@ def labels(self) -> Mapping[str, str]:
 
     @property
     def name(self) -> str:
-        if (name := self._name) is not None:
-            return name
-        return "<unnamed>"
+        return name if (name := self._name) is not None else "<unnamed>"
 
     @staticmethod
     def extend(
@@ -70,14 +72,14 @@ class FnCall:
     annotations: OpAnnotations
 
 
-class StreamObserver(Protocol, Generic[_In_contra]):
-    def on_next(self, log_offset: int, value: _In_contra, /) -> None: ...
+class StreamObserver(Protocol):
+    def on_next(self, log_offset: int, value: object, /) -> None: ...
     def on_close(self, log_offset: int, error: Exception | None, /) -> None: ...
 
 
 @frozen
 class StreamCreate:
-    observer: StreamObserver[Any] | None
+    observer: StreamObserver | None
     dtype: TypeHint[Any]
     annotations: OpAnnotations
 
@@ -123,18 +125,20 @@ class ExternalPromiseComplete:
 
 
 @overload
-def create_op(loop: EventLoop, params: FnCall) -> OpFuture[object]: ...
+def create_op(loop: EventLoop, params: FnCall) -> asyncio.Future[object]: ...
 @overload
-def create_op(loop: EventLoop, params: StreamCreate) -> OpFuture[str]: ...
+def create_op(loop: EventLoop, params: StreamCreate) -> asyncio.Future[str]: ...
 @overload
-def create_op(loop: EventLoop, params: StreamEmit) -> OpFuture[None]: ...
+def create_op(loop: EventLoop, params: StreamEmit) -> asyncio.Future[None]: ...
 @overload
-def create_op(loop: EventLoop, params: StreamClose) -> OpFuture[None]: ...
+def create_op(loop: EventLoop, params: StreamClose) -> asyncio.Future[None]: ...
 @overload
-def create_op(loop: EventLoop, params: Barrier) -> OpFuture[int]: ...
+def create_op(loop: EventLoop, params: Barrier) -> asyncio.Future[int]: ...
 @overload
-def create_op(loop: EventLoop, params: ExternalPromiseCreate) -> OpFuture[object]: ...
+def create_op(loop: EventLoop, params: ExternalPromiseCreate) -> OpFuture: ...
 @overload
-def create_op(loop: EventLoop, params: ExternalPromiseComplete) -> OpFuture[None]: ...
-def create_op(loop: EventLoop, params: Op) -> OpFuture[Any]:
+def create_op(
+    loop: EventLoop, params: ExternalPromiseComplete
+) -> asyncio.Future[None]: ...
+def create_op(loop: EventLoop, params: Op) -> asyncio.Future[Any]:
     return loop.create_op(params, external=asyncio.get_running_loop() is not loop)

src/duron/_core/signal.py

@@ -4,7 +4,7 @@
 import sys
 from asyncio.exceptions import CancelledError
 from collections import deque
-from typing import TYPE_CHECKING, Generic
+from typing import TYPE_CHECKING, Generic, cast
 from typing_extensions import Any, TypeVar, final
 
 from duron._core.ops import (
@@ -131,7 +131,7 @@ async def create_signal(
         loop,
         StreamCreate(
             dtype=dtype,
-            observer=s,
+            observer=cast("Signal[object]", s),
             annotations=annotations,
         ),
     )

src/duron/_core/stream.py

@@ -34,7 +34,7 @@
     from types import TracebackType
 
     from duron._core.context import Context
-    from duron._loop import EventLoop, OpFuture
+    from duron._loop import EventLoop
     from duron.typing import TypeHint
 
     _P = ParamSpec("_P")
@@ -192,7 +192,7 @@ async def create_stream(
         loop,
         StreamCreate(
             dtype=dtype,
-            observer=s,
+            observer=cast("ObserverStream[object, None]", s),
             annotations=annotations,
         ),
     )
@@ -396,7 +396,7 @@ async def __aenter__(self) -> _StreamRun[_U, _T]:
             self._loop,
             StreamCreate(
                 dtype=self._dtype,
-                observer=self._stream,
+                observer=cast("_StreamRun[object, _T]", self._stream),
                 annotations=OpAnnotations.extend(
                     None,
                     name=self._stream.name(),
@@ -457,7 +457,7 @@ def start_worker(self, sink: StreamWriter[_U]) -> asyncio.Future[object]:
                 annotations=OpAnnotations.extend(None, name=self.name()),
             ),
         )
-        self._task = cast("OpFuture[_T]", op)
+        self._task = cast("asyncio.Future[_T]", op)
         return op
 
     async def _worker(self, sink: StreamWriter[_U]) -> _T:

src/duron/_core/stream_manager.py

@@ -21,7 +21,7 @@
 
 @dataclass(slots=True)
 class _StreamInfo:
-    observers: Sequence[StreamObserver[object]]
+    observers: Sequence[StreamObserver]
     dtype: TypeHint[Any]
     labels: Mapping[str, str]
     op_span: OpSpan | None
@@ -33,7 +33,7 @@ class StreamManager:
 
     def __init__(
         self,
-        watchers: list[tuple[dict[str, str], StreamObserver[object]]] | None = None,
+        watchers: list[tuple[dict[str, str], StreamObserver]] | None = None,
     ) -> None:
         self._streams: dict[
             str,
@@ -44,7 +44,7 @@ def __init__(
     def create_stream(
         self,
         stream_id: str,
-        observer: StreamObserver[object] | None,
+        observer: StreamObserver | None,
         dtype: TypeHint[Any],
         labels: Mapping[str, str],
         op_span: OpSpan | None,

src/duron/_loop.py

@@ -9,7 +9,7 @@
 from dataclasses import dataclass
 from hashlib import blake2b
 from heapq import heappop, heappush
-from typing import TYPE_CHECKING, Generic
+from typing import TYPE_CHECKING
 from typing_extensions import (
     Any,
     TypeVar,
@@ -21,14 +21,13 @@
 
 from duron.log import derive_id, random_id
 
-_T = TypeVar("_T")
-
 if TYPE_CHECKING:
     import sys
     from asyncio.futures import Future
     from collections.abc import Callable, Coroutine, Generator
     from contextvars import Context
 
+    _T = TypeVar("_T")
     _Ts = TypeVarTuple("_Ts")
 
     if sys.version_info >= (3, 12):
@@ -37,13 +36,11 @@
         _TaskCompatibleCoro = Generator[Any, None, _T] | Coroutine[Any, Any, _T]
 
 
-logger = logging.getLogger(__name__)
-
-
+logger = logging.getLogger("duron.loop")
 _task_ctx: contextvars.ContextVar[_TaskCtx] = contextvars.ContextVar("duron.task")
 
 
-class OpFuture(asyncio.Future[_T], Generic[_T]):
+class OpFuture(asyncio.Future[object]):
     __slots__: tuple[str, ...] = ("id", "params")
 
     def __init__(
@@ -59,7 +56,7 @@ def __init__(
 
 @dataclass(slots=True)
 class WaitSet:
-    ops: list[OpFuture[object]]
+    ops: list[OpFuture]
     timer: int | None
     event: asyncio.Event
 
@@ -104,7 +101,7 @@ def __init__(self, host: asyncio.AbstractEventLoop) -> None:
         self._exc_handler: (
             Callable[[asyncio.AbstractEventLoop, dict[str, object]], object] | None
         ) = None
-        self._ops: dict[str, OpFuture[object]] = {}
+        self._ops: dict[str, OpFuture] = {}
         self._ctx: _TaskCtx = _TaskCtx(parent_id="")
         self._key: bytes = b""
         self._now_us: int = 0
@@ -261,13 +258,13 @@ def poll_completion(self, task: Future[_T]) -> WaitSet | None:
             if prev_task:
                 tasks._enter_task(self._host, prev_task)  # noqa: SLF001
 
-    def create_op(self, params: object, *, external: bool = False) -> OpFuture[object]:
+    def create_op(self, params: object, *, external: bool = False) -> OpFuture:
         if external:
             id_ = random_id()
             self._event.set()
         else:
             id_ = self.generate_op_id()
-        op_fut: OpFuture[object] = OpFuture(id_, params, self)
+        op_fut = OpFuture(id_, params, self)
         self._ops[id_] = op_fut
         return op_fut