Python abstractmethod with asynccontextmanager

user_4278309 0
Viewed 1966

I have a couple rate limiter classes (one not shown) that I'd like to create an ABC for. The request method is an async context manager. With the code shown below I get

Signature of "request" incompatible with supertype "RateLimiterInterface"

If I try to decorate the abstract method with @asynccontextmanager, I get a typing error:

Argument 1 to "asynccontextmanager" has incompatible type "Callable[[RateLimiterInterface], Coroutine[Any, Any, AsyncIterator[Any]]]"; expected "Callable[..., AsyncIterator[]]"

How can I do this?

class RateLimiterInterface(abc.ABC):
    @abc.abstractmethod
    async def request(self) -> AsyncIterator:
        pass

class LeakyBucketRateLimiter(RateLimiterInterface):
    def __init__(self, max_tokens: Optional[int] = None, rate: float = 60) -> None:
        self.max_tokens = max_tokens
        self.rate = rate
        self._bucket = max_tokens
        self._last_added_at = time.time()

    @contextlib.asynccontextmanager
    async def request(self) -> AsyncIterator:
        if self._bucket is None:
            yield
            return
        while not self._bucket:
            await asyncio.sleep(0)
            self._add_tokens(int((time.time() - self._last_added_at) * self.rate))
        self._bucket -= 1
        yield
        return

    def _add_tokens(self, num_tokens: int) -> None:
        if num_tokens == 0:
            return
        self._bucket += num_tokens
        if self._bucket > self.max_tokens:
            self._bucket = self.max_tokens
        self._last_added_at = time.time()
mypy python-asyncio python-3.x python
2 Answers

I just met same problem with typing and solved it that way:

import abc
import contextlib
import asyncio


class TestAbstract(metaclass=abc.ABCMeta):
    @contextlib.asynccontextmanager
    @abc.abstractmethod
    # Here the trick: you must declare an  asynchronous generator function,
    # not a regular coroutine function so it have to explicitly yield.
    async def foo(self):
        yield


class SubTest(TestAbstract):
    @contextlib.asynccontextmanager
    async def foo(self):
        yield


async def test_coro():
    async with SubTest().foo():
        print('OK')


asyncio.run(test_coro())
user_119264950

I can't reproduce problem. Code below runs fine for me (Python 3.7.1).

Please, test if it runs fine in your environment and if it is alter question to provide minimal reproducible example.

import asyncio
import abc
import contextlib
from typing import AsyncIterator


class Test(abc.ABC):
    @abc.abstractmethod
    async def request(self) -> AsyncIterator:
        pass


class SubTest(Test):
    @contextlib.asynccontextmanager
    async def request(self) -> AsyncIterator:
        await asyncio.sleep(1)
        yield 1
        await asyncio.sleep(1)


async def main():
    obj = SubTest()
    async with obj.request() as res:
        print(res)


asyncio.run(main())
user_11132070
Related