Skip to content

Fix contextlib.asynccontextmanager to work with coroutine functions #10753

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 6 commits into from
Closed

Fix contextlib.asynccontextmanager to work with coroutine functions #10753

wants to merge 6 commits into from

Conversation

sobolevn
Copy link
Member

@sobolevn sobolevn commented Sep 23, 2023 •
edited
Loading

Right now a code from the docs: https://docs.python.org/3/library/contextlib.html#contextlib.asynccontextmanager does not type-check.

Example: https://mypy-play.net/?mypy=latest&python=3.11&gist=f20e230fa7b1af15064b2828dfd7dfd7

We are also required to have AsyncGenerator annotation here, because .athrow() method is used in this function: https://github.com/python/cpython/blob/e8be0c9c5a7c2327b3dd64009f45ee0682322dcb/Lib/contextlib.py#L226 (since 3.8: https://github.com/python/cpython/blob/3.8/Lib/contextlib.py#L189)

Runtime:

>>> from typing import AsyncGenerator
>>> from contextlib import asynccontextmanager
>>> import asyncio
>>> @asynccontextmanager
... async def async_context() -> AsyncGenerator[str, None]:
...     yield 'example'
... 
>>> async def main():
...     async with async_context() as ac:
...         print(ac)
... 
>>> asyncio.run(main())
example

So, I went with the fix that also has some tests.

@github-actions

This comment has been minimized.

Sign in to view
@sobolevn
Copy link
Member Author

Most errors are about AsyncIterator -> AsyncGenerator change. Examples:

@AlexWaygood
Copy link
Member

Most errors are about AsyncIterator -> AsyncGenerator change.

Could you defer that to a separate PR, so we can weigh it up on its own?

@sobolevn
Copy link
Member Author

Technically, this line is not quite correct as well:

typeshed/stdlib/contextlib.pyi

Line 76 in e1b6006

def contextmanager(func: Callable[_P, Iterator[_T_co]]) -> Callable[_P, _GeneratorContextManager[_T_co]]: ...

Because CPython calls self.gen.throw() here: https://github.com/python/cpython/blob/e8be0c9c5a7c2327b3dd64009f45ee0682322dcb/Lib/contextlib.py#L159

@github-actions
Copy link
Contributor

Diff from mypy_primer, showing the effect of this PR on open source code:

prefect (https://github.com/PrefectHQ/prefect)
- src/prefect/client/base.py:37: error: Argument 1 to "asynccontextmanager" has incompatible type "Callable[[Any], AbstractContextManager[None]]"; expected "Callable[[Any], AsyncIterator[<nothing>]]"  [arg-type]
+ src/prefect/client/base.py:37: error: Argument 1 to "asynccontextmanager" has incompatible type "Callable[[Any], AbstractContextManager[None]]"; expected "Callable[[Any], Coroutine[Any, Any, AsyncIterator[<nothing>]]]"  [arg-type]

@AlexWaygood
Copy link
Member

Technically, this line is not quite correct as well:

Indeed. Discussed extensively in #7430 (and, before that, #2773) :)

@sobolevn
Copy link
Member Author

perfect's annotation is not perfect: https://github.com/PrefectHQ/prefect/blob/d739f53f4ae7cbea8ca2f42e42881d4a80481abc/src/prefect/client/base.py#L37-L38

@asynccontextmanager
async def app_lifespan_context(app: FastAPI) -> ContextManager[None]:

So, this is not a regression.

sobolevn added a commit to sobolevn/prefect that referenced this pull request Sep 23, 2023
@sobolevn
Use correct annotation for client.base.app_lifespan_context
4b765bb 
It is `AsyncGenerator[None, None]` and not `ContextManager` :)

Found in python/typeshed#10753
@sobolevn sobolevn mentioned this pull request Sep 23, 2023
Closed
4 tasks
AlexWaygood
AlexWaygood reviewed Sep 23, 2023
View reviewed changes
test_cases/stdlib/check_contextlib.py
Comment on lines +24 to +35
@asynccontextmanager
async def async_context() -> AsyncGenerator[str, None]:
yield "example"


async def async_gen() -> AsyncGenerator[str, None]:
yield "async gen"


@asynccontextmanager
def async_cm_func() -> AsyncGenerator[str, None]:
return async_gen()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All these test cases appear to pass on main

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, mypy already knows how to solve this on main 👍

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, mypy already knows how to solve this on main 👍

Then I'm confused. What's the problem we're trying to fix, if these tests all already pass?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeap, since it was fixed on main, we can drop this fix from typeshed

@AlexWaygood
Copy link
Member

AlexWaygood commented Sep 23, 2023 •
edited
Loading

This fails mypy currently:

from contextlib import asynccontextmanager
from typing import AsyncIterator, AsyncGenerator

@asynccontextmanager
async def db_connection() -> AsyncGenerator[None, None]: ...

But this doesn't:

from contextlib import asynccontextmanager
from typing import AsyncIterator, AsyncGenerator

@asynccontextmanager
async def db_connection() -> AsyncGenerator[None, None]:
    yield None

I think this is because of the rules laid out in https://mypy.readthedocs.io/en/latest/more_types.html#asynchronous-iterators.

https://mypy-play.net/?mypy=latest&python=3.11&gist=1c3911b9a48d2b1ec0e19e0cfee36b16

@sobolevn sobolevn closed this Sep 23, 2023
@AlexWaygood AlexWaygood deleted the fix-asynccontextmanager branch September 23, 2023 13:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@AlexWaygood AlexWaygood AlexWaygood left review comments

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@sobolevn @AlexWaygood