Skip to content

Validate contextmanager return types against the context manager - #378

Open
Muhtasim-Munif-Fahim wants to merge 1 commit into
facebook:mainfrom
Muhtasim-Munif-Fahim:fix/contextmanager-return-type
Open

Validate contextmanager return types against the context manager#378
Muhtasim-Munif-Fahim wants to merge 1 commit into
facebook:mainfrom
Muhtasim-Munif-Fahim:fix/contextmanager-return-type

Conversation

@Muhtasim-Munif-Fahim

Copy link
Copy Markdown

Fixes #193. The issue is from 2020 but still reproduces on current main (555422d) — repro below.

Cause

contextlib.contextmanager and contextlib.asynccontextmanager wrap the decorated generator function with functools.wraps, which copies __annotations__ over verbatim. So for

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

the recorded return type is AsyncGenerator[Foo, None] — what the generator yields — while calling get_cm() actually returns an _AsyncGeneratorContextManager.

_validate_return_type reads that annotation straight off the template, so mocking the function rejects every legitimate return value:

TypeCheckError: type of return must be typing.AsyncGenerator[foo.Foo, NoneType];
got <class 'testslide.core.strict_mock.FooStrictMock'> instead

There is no return value that could satisfy it — the annotation describes something the caller never receives.

Reproduction

Exactly the files from #193; testslide test_foo.py fails before this change and passes after.

foo.py / test_foo.py
# foo.py
from contextlib import asynccontextmanager
from typing import AsyncGenerator


class Foo:
    async def bar(self) -> int:
        return 1

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        return


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


async def call_bar() -> int:
    async with get_cm() as foo:
        return await foo.bar()
# test_foo.py
from testslide.dsl import context
from testslide import StrictMock
from foo import call_bar, Foo

module = "foo"


@context
def foo_test(context):
    context.memoize(
        "foo_mock", lambda self: StrictMock(template=Foo, default_context_manager=True)
    )

    @context.before
    async def mock_get_cm(self):
        self.mock_callable(module, "get_cm").to_return_value(self.foo_mock)
        self.mock_async_callable(self.foo_mock, "bar").to_return_value(3)

    @context.example
    async def test_bar(self):
        self.assertEqual(await call_bar(), 3)

Approach

Detect the decoration by comparing the template against its __wrapped__: contextlib turns a generator function into one that is no longer a generator function, which is what distinguishes it from functools.wraps in general. The expected type then becomes contextlib.AbstractContextManager / AbstractAsyncContextManager.

I used the unparameterized ABCs deliberately. The yielded type is produced later by __aenter__, so it is not present in the value being validated and cannot be checked at this point; claiming AsyncContextManager[Foo] would be checking something we cannot see. Happy to change if you would rather keep the parameter for readability.

The check is redirected, not disabled — a non-context-manager return value is still rejected. fails_for_context_manager_template_given_a_non_context_manager asserts that, so this cannot silently rot into "no validation".

Tests

Adds a ContextManagerTarget fixture (implements both the sync and async protocols) plus decorated factories to tests/sample_module.py, and three examples to the _validate_return_type context. The two positive ones fail without the change with exactly the error from the issue.

Verified on Windows / Python 3.13:

  • lib, mock_callable, mock_async_callable, strict_mock, mock_constructor, patch_attribute: 2865 successful, 4 failed, 62 skipped — identical baseline of 4 pre-existing failures before and after; the 3 new successes are mine.
  • unittest suites: 141 tests, only cli_unittest erroring, which is import pty being Unix-only, and reproduces on a clean checkout.
  • ruff check back to the same 22 pre-existing findings; flake8 --select=F,C90 clean. I left the pre-existing ruff format deviation in sample_module.py (a lambda in Target.__init__) untouched so the diff stays on topic.

🤖 Generated with Claude Code

`contextlib.contextmanager` and `contextlib.asynccontextmanager` wrap the
decorated generator function with `functools.wraps`, which copies
`__annotations__` over verbatim. The recorded return type therefore describes
what the *generator* yields, while calling the decorated function actually
returns a context manager.

`_validate_return_type` read that annotation straight off the template, so
mocking such a function rejected every legitimate return value:

    TypeCheckError: type of return must be
    typing.AsyncGenerator[foo.Foo, NoneType]; got
    <class 'testslide.core.strict_mock.FooStrictMock'> instead

Detect the decoration by comparing the template against its `__wrapped__`:
contextlib turns a generator function into one that is no longer a generator
function, which distinguishes it from `functools.wraps` in general. The
expected type then becomes `contextlib.AbstractContextManager` or
`AbstractAsyncContextManager`.

The check is redirected, not disabled: a non-context-manager return value is
still rejected, which the third new test asserts.

Adds a `ContextManagerTarget` fixture plus decorated sync and async factories
to tests/sample_module.py, and three examples to the `_validate_return_type`
context. The two positive ones fail without this change.

Fixes facebook#193

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@meta-cla

meta-cla Bot commented Aug 13, 2026

Copy link
Copy Markdown

Hi @Muhtasim-Munif-Fahim!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Type error when using contextlib.asynccontextmanager

1 participant