Apply coding best practices: fix mypy strict, remove future annotations, improve async concurrency, hoist imports

Agent-Logs-Url: https://github.com/codeSamuraii/pyfuse/sessions/b525c00f-293b-4d49-8924-ad15d185248a

Co-authored-by: codeSamuraii <17270548+codeSamuraii@users.noreply.github.com>

Author: Copilot

pyfuse/__main__.py

@@ -15,6 +15,7 @@
 import importlib
 import logging
 import os
+import shutil
 import signal
 import sys
 from collections.abc import Callable
@@ -385,8 +386,6 @@ def _cmd_sandbox(args: argparse.Namespace) -> None:
 
 async def _sandbox_status() -> None:
     """Print the current Docker sandbox status."""
-    import shutil
-
     if shutil.which("docker") is not None:
         from pyfuse.worker.sandbox.docker import (
             _container_exists, _container_running, _image_exists,
@@ -410,7 +409,6 @@ async def _sandbox_status() -> None:
 
 async def _docker_setup() -> None:
     """Build the Docker sandbox image."""
-    import shutil
     if shutil.which("docker") is None:
         print(
             "Error: 'docker' command not found.\n"
@@ -432,7 +430,6 @@ async def _docker_setup() -> None:
 
 async def _docker_teardown() -> None:
     """Stop and remove the Docker sandbox container and image."""
-    import shutil
     if shutil.which("docker") is None:
         print("Docker is not installed, nothing to tear down.", file=sys.stderr)
         return

pyfuse/graph/tracing.py

@@ -1,3 +1,4 @@
+import asyncio
 import builtins
 import contextvars
 import functools
@@ -7,7 +8,7 @@
 import sys
 import sysconfig
 import threading
-from collections.abc import AsyncGenerator, Callable, Generator
+from collections.abc import AsyncGenerator, Awaitable, Callable, Generator
 from pathlib import Path
 from typing import Any, ParamSpec, TypeVar, cast
 
@@ -54,8 +55,15 @@ def _make_map_method(
     """Create the ``.map()`` async method for batch submission and collection."""
 
     async def map(args_list: list[tuple[object, ...]], **kwargs: object) -> list[object]:
-        results = [await start_method(*args, **kwargs) for args in args_list]  # type: ignore[misc]
-        return [await r for r in results]
+        coros: list[Awaitable[object]] = [
+            cast(Awaitable[object], start_method(*args, **kwargs))
+            for args in args_list
+        ]
+        results = await asyncio.gather(*coros)
+        awaitables: list[Awaitable[object]] = [
+            cast(Awaitable[object], r) for r in results
+        ]
+        return list(await asyncio.gather(*awaitables))
 
     return map
 

pyfuse/worker/backends/redis.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 try:
-    import redis.asyncio as _redis
+    import redis.asyncio as _redis  # type: ignore[import-not-found]
 except ImportError:
     raise ImportError(
         "redis package is required for RedisBackend. "

pyfuse/worker/remote.py

@@ -9,7 +9,7 @@
 import sys
 import time
 from collections.abc import Awaitable, Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pyfuse.core.progress import _progress_callback
 from pyfuse.core.task import Task
@@ -18,6 +18,9 @@
 from pyfuse.worker.result import Result, ResultEnvelope
 from pyfuse.worker.worker import Worker
 
+if TYPE_CHECKING:
+    from pyfuse.worker.sandbox import DockerSandbox
+
 logger = logging.getLogger(__name__)
 
 _active_backend: Backend | None = None
@@ -505,6 +508,7 @@ async def serve(
     # Boot the sandbox container before accepting tasks so the first
     # execution doesn't pay the startup cost.
     if worker.sandboxed:
+        assert worker._sandbox is not None
         await worker._sandbox.start()
 
     logger.info("Listening for tasks \u2014 Ctrl+C to stop.")
@@ -513,6 +517,7 @@ async def serve(
         await _worker_loop(worker, backend, concurrency)
     finally:
         if worker.sandboxed:
+            assert worker._sandbox is not None
             await worker._sandbox.stop()
         await disconnect()
         logger.info("Worker stopped.")

pyfuse/worker/sandbox/_protocol.py

@@ -23,7 +23,7 @@ def encode(obj: dict[str, Any]) -> bytes:
 
 def decode_header(data: bytes) -> int:
     """Return the payload length from a 4-byte header."""
-    (length,) = _HEADER.unpack(data)
+    length: int = _HEADER.unpack(data)[0]
     return length
 
 

pyfuse/worker/sandbox/docker.py

@@ -16,8 +16,6 @@
 CI to avoid a cold-start build).
 """
 
-from __future__ import annotations
-
 import asyncio
 import logging
 import os
@@ -27,6 +25,7 @@
 from typing import Any
 
 from pyfuse.core.errors import WorkerError
+from pyfuse.core.progress import _progress_callback
 from pyfuse.worker.sandbox._protocol import async_recv, async_send
 
 logger = logging.getLogger(__name__)
@@ -113,7 +112,6 @@ async def execute(
 
         # Pick up the host-side progress callback (set by _handle_task)
         # so we can forward progress messages from the container.
-        from pyfuse.core.progress import _progress_callback
         progress_cb = _progress_callback.get(None)
 
         try:

pyfuse/worker/sandbox/guest_agent.py

@@ -38,8 +38,6 @@
     python guest_agent.py [--host 0.0.0.0] [--port 9749]
 """
 
-from __future__ import annotations
-
 import argparse
 import asyncio
 import contextvars

pyfuse/worker/worker.py

@@ -135,6 +135,7 @@ async def run(self, task: Task) -> Any:
         logger.debug("Executing %s (cache key: %s)", task.function_name, cached.subgraph_key)
 
         if self.sandboxed:
+            assert self._sandbox is not None
             # Determine owner_class for method dispatch inside the sandbox
             store = Store.from_json(task.graph_json)
             target_qname, nodes = store.collect(task.function_name)