From ba17daa7f705254961bd2b82c153b0d16fa0d35c Mon Sep 17 00:00:00 2001 From: Muhtasim-Munif-Fahim Date: Thu, 13 Aug 2026 21:55:07 +0600 Subject: [PATCH] fix: validate contextmanager return types against the context manager `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 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 #193 Co-Authored-By: Claude Opus 5 --- tests/lib_testslide.py | 26 ++++++++++++++++++++++++++ tests/sample_module.py | 37 ++++++++++++++++++++++++++++++++++++- testslide/core/lib.py | 31 +++++++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 1 deletion(-) diff --git a/tests/lib_testslide.py b/tests/lib_testslide.py index ec60bb31..f76ffcb8 100644 --- a/tests/lib_testslide.py +++ b/tests/lib_testslide.py @@ -376,6 +376,32 @@ def passes_for_valid_forward_reference(self): Foo.get_maybe_foo, Foo(), self.caller_frame_info ) + @context.example + def passes_for_context_manager_template(self): + """ + contextlib.contextmanager copies __annotations__ from the generator + function it wraps, so the recorded return type describes what the + generator yields rather than the context manager that calling it + actually returns. + https://github.com/facebook/TestSlide/issues/193 + """ + self.callable_template = sample_module.test_function_returns_context_manager + self.assert_passes(StrictMock(template=sample_module.ContextManagerTarget)) + + @context.example + def passes_for_async_context_manager_template(self): + """Same as above, for contextlib.asynccontextmanager.""" + self.callable_template = ( + sample_module.test_function_returns_async_context_manager + ) + self.assert_passes(StrictMock(template=sample_module.ContextManagerTarget)) + + @context.example + def fails_for_context_manager_template_given_a_non_context_manager(self): + """The check is redirected, not disabled.""" + self.callable_template = sample_module.test_function_returns_context_manager + self.assert_fails(42) + @context.example def fails_for_valid_forward_reference_but_bad_type_passed(self): with self.assertRaisesRegex( diff --git a/tests/sample_module.py b/tests/sample_module.py index 3a7d787a..4d1413d1 100644 --- a/tests/sample_module.py +++ b/tests/sample_module.py @@ -3,9 +3,12 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from collections.abc import Awaitable, Coroutine +from collections.abc import AsyncGenerator, Awaitable, Coroutine, Generator +from contextlib import asynccontextmanager, contextmanager from typing import Any, Union +from typing_extensions import Self + attribute = "value" typedattr: str = "bruh" @@ -175,3 +178,35 @@ def test_union(arg: UnionArgType) -> None: def test_tuple(arg: TupleArgType) -> None: pass + + +class ContextManagerTarget: + "This class is used by some unit tests only" + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: + return None + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: + return None + + +@contextmanager +def test_function_returns_context_manager() -> Generator[ + "ContextManagerTarget", None, None +]: + "This function is used by some unit tests only" + yield ContextManagerTarget() + + +@asynccontextmanager +async def test_function_returns_async_context_manager() -> AsyncGenerator[ + "ContextManagerTarget", None +]: + "This function is used by some unit tests only" + yield ContextManagerTarget() diff --git a/testslide/core/lib.py b/testslide/core/lib.py index 57fe61df..00a9a9fe 100644 --- a/testslide/core/lib.py +++ b/testslide/core/lib.py @@ -5,6 +5,7 @@ # pyre-unsafe import collections.abc as abc +import contextlib import functools import inspect import os @@ -341,6 +342,33 @@ def _is_wrapped_for_signature_and_type_validation(value: Callable) -> bool: return getattr(value, "__is_testslide_type_validation_wrapping", False) +def _get_contextmanager_return_type(template: Any) -> Any: + """Return the context manager type produced by a contextlib-decorated function. + + ``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 (eg ``AsyncGenerator[Foo, None]``), + while calling the decorated function actually returns a context manager. + Validating a mocked return value against the generator type rejects every + legitimate value. + + Returns ``None`` when ``template`` is not such a decorated function. + """ + wrapped = getattr(template, "__wrapped__", None) + if wrapped is None: + return None + # The decorator turns a generator function into one that is no longer a + # generator function, which is what distinguishes it from @wraps in general. + if inspect.isasyncgenfunction(wrapped) and not inspect.isasyncgenfunction(template): + return contextlib.AbstractAsyncContextManager + if inspect.isgeneratorfunction(wrapped) and not inspect.isgeneratorfunction( + template + ): + return contextlib.AbstractContextManager + return None + + def _validate_return_type( template: Mock | Callable, value: Any, @@ -352,6 +380,9 @@ def _validate_return_type( except TypeError: return expected_type = argspec.annotations.get("return") + contextmanager_type = _get_contextmanager_return_type(template) + if contextmanager_type is not None: + expected_type = contextmanager_type if expected_type: if unwrap_template_awaitable: type_origin = get_origin(expected_type)