2025-10-16

[MtkN1]

Today I Learned ...

[MtkN1]

非同期バッチ実行支援モジュールのスケッチ

非同期実行フレームワーク非依存で、型安全にバッチ実行を支援するモジュール:

example_batch.py

__all__ = [
    "AsyncBatch",
    "Future",
]

import functools
from collections.abc import Awaitable, Callable


class _MissingType: ...


_MISSING = _MissingType()


class Future[T]:
    def __init__(self) -> None:
        self._result: T | _MissingType = _MISSING

    def set_result(self, result: T, /) -> None:
        self._result = result

    def result(self) -> T:
        if isinstance(self._result, _MissingType):
            raise RuntimeError

        return self._result


class ManyFutures[T]:
    def __init__(self) -> None:
        self._futures: list[Future[T]] = []

    def add_future(self, future: Future[T], /) -> None:
        self._futures.append(future)

    def results(self) -> tuple[T, ...]:
        return tuple(f.result() for f in self._futures)


class AsyncBatch:
    def __init__(self) -> None:
        self._calls: list[Callable[[], Awaitable[object]]] = []

    @staticmethod
    async def _wraps[*P, R](
        future: Future[R],
        async_fn: Callable[[*P], Awaitable[R]],
        /,
        *args: *P,
    ) -> Future[R]:
        result = await async_fn(*args)
        future.set_result(result)
        return future

    def add[*P, R](
        self,
        async_fn: Callable[[*P], Awaitable[R]],
        /,
        *args: *P,
    ) -> Future[R]:
        future = Future[R]()
        wrapped = functools.partial(self._wraps, future, async_fn, *args)
        self._calls.append(wrapped)
        return future

    def add_many[*P, R](
        self,
        async_fn: Callable[[*P], Awaitable[R]],
        /,
        parameters: tuple[tuple[*P], ...],
    ) -> ManyFutures[R]:
        futures = ManyFutures[R]()

        for args in parameters:
            future = self.add(async_fn, *args)
            futures.add_future(future)

        return futures

    @property
    def calls(self) -> tuple[Callable[[], Awaitable[object]], ...]:
        return tuple(self._calls)

example_batch_test.py

import pytest
import trio

import example_batch


async def fetch_int(num: int = 42, /) -> int:
    return num


async def fetch_str(msg: str = "hello", /) -> str:
    return msg


@pytest.mark.trio
async def test_async_batch() -> None:
    async with trio.open_nursery() as nursery:
        batch = example_batch.AsyncBatch()

        future_int = batch.add(fetch_int)
        future_str = batch.add(fetch_str)

        for call in batch.calls:
            nursery.start_soon(call)

    assert future_int.result() == 42
    assert future_str.result() == "hello"


@pytest.mark.trio
async def test_async_batch_many() -> None:
    async with trio.open_nursery() as nursery:
        batch = example_batch.AsyncBatch()

        futures_int = batch.add_many(fetch_int, tuple((i,) for i in range(3)))
        futures_str = batch.add_many(fetch_str, tuple((s,) for s in ("hello", "world")))

        for call in batch.calls:
            nursery.start_soon(call)

    assert futures_int.results() == (0, 1, 2)
    assert futures_str.results() == ("hello", "world")