From dc91faf357e949507e3b2a136a1dccd29eac8c75 Mon Sep 17 00:00:00 2001 From: Conrad Date: Fri, 20 Feb 2026 22:18:10 -0500 Subject: [PATCH 1/3] build: Add pyright dev dependency and configuration Add pyright to dev dependencies and configure it in pyproject.toml with appropriate include/exclude paths. Excludes protobuf generated files, test directories, and cli.py scratch module from analysis. --- wool/pyproject.toml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/wool/pyproject.toml b/wool/pyproject.toml index 341a508..ecd5bd6 100644 --- a/wool/pyproject.toml +++ b/wool/pyproject.toml @@ -41,6 +41,7 @@ dev = [ "cryptography", "debugpy", "hypothesis", + "pyright", "pytest", "pytest-asyncio", "pytest-cov", @@ -96,3 +97,17 @@ docstring-code-format = true combine-as-imports = false force-single-line = true known-first-party = ["wool"] + +[tool.pyright] +venvPath = "." +venv = ".venv" +include = ["src"] +exclude = [ + "**/__pycache__", + "build", + "dist", + "tests/*", + "src/wool/runtime/protobuf/*pb2*", +] +reportMissingImports = true +reportMissingTypeStubs = false From f809636e8fba724900c8138891f30d69e5d27b4e Mon Sep 17 00:00:00 2001 From: Conrad Date: Fri, 20 Feb 2026 22:18:19 -0500 Subject: [PATCH 2/3] ci: Add pyright linter job to run-tests workflow --- .github/workflows/run-tests.yaml | 34 +++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml index ec47aa1..e214af2 100644 --- a/.github/workflows/run-tests.yaml +++ b/.github/workflows/run-tests.yaml @@ -8,6 +8,28 @@ on: - release jobs: + lint: + name: Lint / pyright + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install uv and prepare python + uses: astral-sh/setup-uv@v5 + with: + python-version: '3.13' + + - name: Install packages + run: | + .github/scripts/install-python-packages.sh + uv pip install -e './wool[dev]' + + - name: Run pyright + run: | + cd wool + uv run pyright + run-tests: name: Namespace ${{ matrix.namespace }} / Python ${{ matrix.python-version }} runs-on: ubuntu-latest @@ -15,24 +37,30 @@ jobs: matrix: namespace: - 'wool' - python-version: + python-version: - '3.11' - '3.12' - '3.13' steps: - name: Checkout code uses: actions/checkout@v4 + - name: Install uv and prepare python uses: astral-sh/setup-uv@v5 with: python-version: ${{ matrix.python-version }} + - name: Install packages - run: .github/scripts/install-python-packages.sh - - name: Run tests env: NAMESPACE: ${{ matrix.namespace }} run: | + .github/scripts/install-python-packages.sh uv pip install -e './${{ env.NAMESPACE }}[dev]' uv pip freeze + + - name: Run tests + env: + NAMESPACE: ${{ matrix.namespace }} + run: | cd ${{ env.NAMESPACE }} uv run pytest From bb476b28ec8a2702ed02c29a9e023eb225e9b525 Mon Sep 17 00:00:00 2001 From: Conrad Date: Fri, 20 Feb 2026 22:19:32 -0500 Subject: [PATCH 3/3] fix: Resolve type errors surfaced by pyright Correct return type annotations and narrowing issues caught by pyright across the runtime subpackages: - roundrobin: Use AsyncGenerator instead of AsyncIterator for dispatch - task: Add casts for generic W return type and runtime assertions for callable type narrowing - connection: Annotate _execute return type and cast the generator in dispatch to satisfy AsyncGenerator variance - pool: Narrow overload 1 discovery parameter to None to prevent overlap with the discovery-required overload; add third overload for hybrid pools that both spawn local workers and discover remote workers --- .../wool/runtime/loadbalancer/roundrobin.py | 4 +-- wool/src/wool/runtime/routine/task.py | 8 ++++-- wool/src/wool/runtime/worker/connection.py | 10 +++++-- wool/src/wool/runtime/worker/pool.py | 28 +++++++++++++++---- 4 files changed, 38 insertions(+), 12 deletions(-) diff --git a/wool/src/wool/runtime/loadbalancer/roundrobin.py b/wool/src/wool/runtime/loadbalancer/roundrobin.py index d9e2eef..3cdde39 100644 --- a/wool/src/wool/runtime/loadbalancer/roundrobin.py +++ b/wool/src/wool/runtime/loadbalancer/roundrobin.py @@ -4,7 +4,7 @@ import logging from asyncio import Lock from typing import TYPE_CHECKING -from typing import AsyncIterator +from typing import AsyncGenerator from typing import Final from wool.runtime.worker.connection import TransientRpcError @@ -43,7 +43,7 @@ async def dispatch( *, context: LoadBalancerContextLike, timeout: float | None = None, - ) -> AsyncIterator: + ) -> AsyncGenerator: """Dispatch a task to the next available worker. :param task: diff --git a/wool/src/wool/runtime/routine/task.py b/wool/src/wool/runtime/routine/task.py index 45a1a7b..a199423 100644 --- a/wool/src/wool/runtime/routine/task.py +++ b/wool/src/wool/runtime/routine/task.py @@ -12,6 +12,7 @@ from inspect import isasyncgenfunction from inspect import iscoroutinefunction from types import TracebackType +from typing import Any from typing import AsyncGenerator from typing import ContextManager from typing import Coroutine @@ -23,6 +24,7 @@ from typing import Tuple from typing import TypeAlias from typing import TypeVar +from typing import cast from typing import overload from uuid import UUID @@ -236,9 +238,9 @@ def to_protobuf(self) -> pb.task.Task: def dispatch(self) -> W: if isasyncgenfunction(self.callable): - return self._stream() + return cast(W, self._stream()) elif iscoroutinefunction(self.callable): - return self._run() + return cast(W, self._run()) else: raise ValueError("Expected routine to be coroutine or async generator") @@ -251,6 +253,7 @@ async def _run(self): :raises RuntimeError: If no proxy pool is available for task execution. """ + assert iscoroutinefunction(self.callable), "Expected coroutine function" proxy_pool = wool.__proxy_pool__.get() if not proxy_pool: raise RuntimeError("No proxy pool available for task execution") @@ -274,6 +277,7 @@ async def _stream(self): :raises RuntimeError: If no proxy pool is available for task execution. """ + assert isasyncgenfunction(self.callable), "Expected async generator function" proxy_pool = wool.__proxy_pool__.get() if not proxy_pool: raise RuntimeError("No proxy pool available for task execution") diff --git a/wool/src/wool/runtime/worker/connection.py b/wool/src/wool/runtime/worker/connection.py index 44d015b..dea15a3 100644 --- a/wool/src/wool/runtime/worker/connection.py +++ b/wool/src/wool/runtime/worker/connection.py @@ -2,11 +2,13 @@ import asyncio from dataclasses import dataclass +from typing import AsyncGenerator from typing import AsyncIterator from typing import Final from typing import Generic from typing import TypeAlias from typing import TypeVar +from typing import cast import cloudpickle import grpc.aio @@ -288,7 +290,7 @@ async def dispatch( task: Task, *, timeout: float | None = None, - ) -> AsyncIterator[pb.task.Result]: + ) -> AsyncGenerator[pb.task.Result, None]: """Dispatch a task to the remote worker for execution. Sends the task to the worker via gRPC, waits for acknowledgment, @@ -338,7 +340,7 @@ async def dispatch( raise await _channel_pool.release(self._key) - return gen + return cast(AsyncGenerator[pb.task.Result, None], gen) async def close(self): """Close the connection and release all pooled resources. @@ -375,7 +377,9 @@ async def _dispatch(self, ch, task, timeout): raise return call - async def _execute(self, call): + async def _execute( + self, call: _DispatchCall + ) -> AsyncGenerator[pb.task.Result | None, None]: ch = await _channel_pool.acquire(self._key) try: yield # Priming yield — signals dispatch() that ref is held diff --git a/wool/src/wool/runtime/worker/pool.py b/wool/src/wool/runtime/worker/pool.py index 2be85c1..1c44be3 100644 --- a/wool/src/wool/runtime/worker/pool.py +++ b/wool/src/wool/runtime/worker/pool.py @@ -160,15 +160,15 @@ def __init__( *tags: str, size: int = 0, worker: WorkerFactory = LocalWorker, - discovery: DiscoveryLike | Factory[DiscoveryLike] | None = None, + discovery: None = None, loadbalancer: ( LoadBalancerLike | Factory[LoadBalancerLike] ) = RoundRobinLoadBalancer, credentials: WorkerCredentials | None | UndefinedType = Undefined, ): """ - Create an ephemeral pool of workers, spawning the specified quantity of workers - using the specified worker factory. + Create an ephemeral pool of workers, spawning the specified + quantity of workers using the specified worker factory. """ ... @@ -183,8 +183,26 @@ def __init__( credentials: WorkerCredentials | None | UndefinedType = Undefined, ): """ - Connect to an existing pool of workers discovered by the specified discovery - protocol. + Connect to an existing pool of workers discovered by the + specified discovery protocol. + """ + ... + + @overload + def __init__( + self, + *tags: str, + size: int = 0, + worker: WorkerFactory = LocalWorker, + discovery: DiscoveryLike | Factory[DiscoveryLike], + loadbalancer: ( + LoadBalancerLike | Factory[LoadBalancerLike] + ) = RoundRobinLoadBalancer, + credentials: WorkerCredentials | None | UndefinedType = Undefined, + ): + """ + Create a hybrid pool that spawns local workers and discovers + remote workers through the specified discovery protocol. """ ...