From 934fd37d2bddf075d27aa7ee1c9f912f139f4ffe Mon Sep 17 00:00:00 2001 From: Simon Lamon <32477463+silamon@users.noreply.github.com> Date: Tue, 5 Aug 2025 16:41:37 +0200 Subject: [PATCH] [ty] Diagnostics for async context managers (#19704) ## Summary Implements diagnostics for async context managers. Fixes https://github.com/astral-sh/ty/issues/918. ## Test Plan Mdtests have been added. --- ...ccidental_use_of_as…_(5b8c1b4d846bc544).snap | 41 +++++ ...idental_use_of_no…_(b07503f9b773ea61).snap | 60 +------ .../resources/mdtest/with/async.md | 165 ++++++++++++++++++ .../resources/mdtest/with/sync.md | 6 +- crates/ty_python_semantic/src/types.rs | 150 +++++++++++----- crates/ty_python_semantic/src/types/class.rs | 6 +- crates/ty_python_semantic/src/types/infer.rs | 10 +- .../ty_python_semantic/src/types/unpacker.rs | 24 ++- 8 files changed, 341 insertions(+), 121 deletions(-) create mode 100644 crates/ty_python_semantic/resources/mdtest/snapshots/async.md_-_Async_with_statement…_-_Accidental_use_of_as…_(5b8c1b4d846bc544).snap diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/async.md_-_Async_with_statement…_-_Accidental_use_of_as…_(5b8c1b4d846bc544).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/async.md_-_Async_with_statement…_-_Accidental_use_of_as…_(5b8c1b4d846bc544).snap new file mode 100644 index 0000000000..2b6dddc52f --- /dev/null +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/async.md_-_Async_with_statement…_-_Accidental_use_of_as…_(5b8c1b4d846bc544).snap @@ -0,0 +1,41 @@ +--- +source: crates/ty_test/src/lib.rs +expression: snapshot +--- +--- +mdtest name: async.md - Async with statements - Accidental use of async `async with` +mdtest path: crates/ty_python_semantic/resources/mdtest/with/async.md +--- + +# Python source files + +## mdtest_snippet.py + +``` +1 | class Manager: +2 | def __enter__(self): ... +3 | def __exit__(self, *args): ... +4 | +5 | async def main(): +6 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aenter__` and `__aexit__`" +7 | async with Manager(): +8 | ... +``` + +# Diagnostics + +``` +error[invalid-context-manager]: Object of type `Manager` cannot be used with `async with` because it does not implement `__aenter__` and `__aexit__` + --> src/mdtest_snippet.py:7:16 + | +5 | async def main(): +6 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aent… +7 | async with Manager(): + | ^^^^^^^^^ +8 | ... + | +info: Objects of type `Manager` can be used as sync context managers +info: Consider using `with` here +info: rule `invalid-context-manager` is enabled by default + +``` diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/sync.md_-_With_statements_-_Accidental_use_of_no…_(b07503f9b773ea61).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/sync.md_-_With_statements_-_Accidental_use_of_no…_(b07503f9b773ea61).snap index 175e8ca13f..25a2eeb705 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/sync.md_-_With_statements_-_Accidental_use_of_no…_(b07503f9b773ea61).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/sync.md_-_With_statements_-_Accidental_use_of_no…_(b07503f9b773ea61).snap @@ -12,27 +12,13 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/with/sync.md ## mdtest_snippet.py ``` - 1 | class Manager: - 2 | async def __aenter__(self): ... - 3 | async def __aexit__(self, *args): ... - 4 | - 5 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and `__exit__`" - 6 | with Manager(): - 7 | ... - 8 | class Manager: - 9 | async def __aenter__(self): ... -10 | async def __aexit__(self, typ: str, exc, traceback): ... -11 | -12 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and `__exit__`" -13 | with Manager(): -14 | ... -15 | class Manager: -16 | async def __aenter__(self, wrong_extra_arg): ... -17 | async def __aexit__(self, typ, exc, traceback, wrong_extra_arg): ... -18 | -19 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and `__exit__`" -20 | with Manager(): -21 | ... +1 | class Manager: +2 | async def __aenter__(self): ... +3 | async def __aexit__(self, *args): ... +4 | +5 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and `__exit__`" +6 | with Manager(): +7 | ... ``` # Diagnostics @@ -45,41 +31,9 @@ error[invalid-context-manager]: Object of type `Manager` cannot be used with `wi 6 | with Manager(): | ^^^^^^^^^ 7 | ... -8 | class Manager: | info: Objects of type `Manager` can be used as async context managers info: Consider using `async with` here info: rule `invalid-context-manager` is enabled by default ``` - -``` -error[invalid-context-manager]: Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and `__exit__` - --> src/mdtest_snippet.py:13:6 - | -12 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and … -13 | with Manager(): - | ^^^^^^^^^ -14 | ... -15 | class Manager: - | -info: Objects of type `Manager` can be used as async context managers -info: Consider using `async with` here -info: rule `invalid-context-manager` is enabled by default - -``` - -``` -error[invalid-context-manager]: Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and `__exit__` - --> src/mdtest_snippet.py:20:6 - | -19 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and … -20 | with Manager(): - | ^^^^^^^^^ -21 | ... - | -info: Objects of type `Manager` can be used as async context managers -info: Consider using `async with` here -info: rule `invalid-context-manager` is enabled by default - -``` diff --git a/crates/ty_python_semantic/resources/mdtest/with/async.md b/crates/ty_python_semantic/resources/mdtest/with/async.md index 1afa89984b..2d907fa381 100644 --- a/crates/ty_python_semantic/resources/mdtest/with/async.md +++ b/crates/ty_python_semantic/resources/mdtest/with/async.md @@ -35,6 +35,171 @@ async def test(): reveal_type(y) # revealed: str ``` +## Context manager without an `__aenter__` or `__aexit__` method + +```py +class Manager: ... + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aenter__` and `__aexit__`" + async with Manager(): + ... +``` + +## Context manager without an `__aenter__` method + +```py +class Manager: + async def __aexit__(self, exc_tpe, exc_value, traceback): ... + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aenter__`" + async with Manager(): + ... +``` + +## Context manager without an `__aexit__` method + +```py +class Manager: + async def __aenter__(self): ... + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aexit__`" + async with Manager(): + ... +``` + +## Context manager with non-callable `__aenter__` attribute + +```py +class Manager: + __aenter__: int = 42 + + async def __aexit__(self, exc_tpe, exc_value, traceback): ... + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not correctly implement `__aenter__`" + async with Manager(): + ... +``` + +## Context manager with non-callable `__aexit__` attribute + +```py +from typing_extensions import Self + +class Manager: + def __aenter__(self) -> Self: + return self + __aexit__: int = 32 + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not correctly implement `__aexit__`" + async with Manager(): + ... +``` + +## Context expression with possibly-unbound union variants + +```py +async def _(flag: bool): + class Manager1: + def __aenter__(self) -> str: + return "foo" + + def __aexit__(self, exc_type, exc_value, traceback): ... + + class NotAContextManager: ... + context_expr = Manager1() if flag else NotAContextManager() + + # error: [invalid-context-manager] "Object of type `Manager1 | NotAContextManager` cannot be used with `async with` because the methods `__aenter__` and `__aexit__` are possibly unbound" + async with context_expr as f: + reveal_type(f) # revealed: str +``` + +## Context expression with "sometimes" callable `__aenter__` method + +```py +async def _(flag: bool): + class Manager: + if flag: + async def __aenter__(self) -> str: + return "abcd" + + async def __exit__(self, *args): ... + + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because the method `__aenter__` is possibly unbound" + async with Manager() as f: + reveal_type(f) # revealed: CoroutineType[Any, Any, str] +``` + +## Invalid `__aenter__` signature + +```py +class Manager: + async def __aenter__() -> str: + return "foo" + + async def __aexit__(self, exc_type, exc_value, traceback): ... + +async def main(): + context_expr = Manager() + + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not correctly implement `__aenter__`" + async with context_expr as f: + reveal_type(f) # revealed: CoroutineType[Any, Any, str] +``` + +## Accidental use of async `async with` + + + +If a asynchronous `async with` statement is used on a type with `__enter__` and `__exit__`, we show +a diagnostic hint that the user might have intended to use `with` instead. + +```py +class Manager: + def __enter__(self): ... + def __exit__(self, *args): ... + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aenter__` and `__aexit__`" + async with Manager(): + ... +``` + +## Incorrect signatures + +The sub-diagnostic is also provided if the signatures of `__enter__` and `__exit__` do not match the +expected signatures for a context manager: + +```py +class Manager: + def __enter__(self): ... + def __exit__(self, typ: str, exc, traceback): ... + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aenter__` and `__aexit__`" + async with Manager(): + ... +``` + +## Incorrect number of arguments + +Similarly, we also show the hint if the functions have the wrong number of arguments: + +```py +class Manager: + def __enter__(self, wrong_extra_arg): ... + def __exit__(self, typ, exc, traceback, wrong_extra_arg): ... + +async def main(): + # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `async with` because it does not implement `__aenter__` and `__aexit__`" + async with Manager(): + ... +``` + ## `@asynccontextmanager` ```py diff --git a/crates/ty_python_semantic/resources/mdtest/with/sync.md b/crates/ty_python_semantic/resources/mdtest/with/sync.md index 89d0b281da..1dd8d0ea70 100644 --- a/crates/ty_python_semantic/resources/mdtest/with/sync.md +++ b/crates/ty_python_semantic/resources/mdtest/with/sync.md @@ -155,7 +155,7 @@ with context_expr as f: If a synchronous `with` statement is used on a type with `__aenter__` and `__aexit__`, we show a -diagnostic hint that the user might have intended to use `asnyc with` instead. +diagnostic hint that the user might have intended to use `async with` instead. ```py class Manager: @@ -167,6 +167,8 @@ with Manager(): ... ``` +## Incorrect signatures + The sub-diagnostic is also provided if the signatures of `__aenter__` and `__aexit__` do not match the expected signatures for a context manager: @@ -180,6 +182,8 @@ with Manager(): ... ``` +## Incorrect number of arguments + Similarly, we also show the hint if the functions have the wrong number of arguments: ```py diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index c171013e20..6f85606830 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -4880,53 +4880,78 @@ impl<'db> Type<'db> { /// Returns the type bound from a context manager with type `self`. /// /// This method should only be used outside of type checking because it omits any errors. - /// For type checking, use [`try_enter`](Self::try_enter) instead. + /// For type checking, use [`try_enter_with_mode`](Self::try_enter_with_mode) instead. fn enter(self, db: &'db dyn Db) -> Type<'db> { - self.try_enter(db) + self.try_enter_with_mode(db, EvaluationMode::Sync) + .unwrap_or_else(|err| err.fallback_enter_type(db)) + } + + /// Returns the type bound from a context manager with type `self`. + /// + /// This method should only be used outside of type checking because it omits any errors. + /// For type checking, use [`try_enter_with_mode`](Self::try_enter_with_mode) instead. + fn aenter(self, db: &'db dyn Db) -> Type<'db> { + self.try_enter_with_mode(db, EvaluationMode::Async) .unwrap_or_else(|err| err.fallback_enter_type(db)) } /// Given the type of an object that is used as a context manager (i.e. in a `with` statement), - /// return the return type of its `__enter__` method, which is bound to any potential targets. + /// return the return type of its `__enter__` or `__aenter__` method, which is bound to any potential targets. /// /// E.g., for the following `with` statement, given the type of `x`, infer the type of `y`: /// ```python /// with x as y: /// pass /// ``` - fn try_enter(self, db: &'db dyn Db) -> Result, ContextManagerError<'db>> { - let enter = self.try_call_dunder(db, "__enter__", CallArguments::none()); + fn try_enter_with_mode( + self, + db: &'db dyn Db, + mode: EvaluationMode, + ) -> Result, ContextManagerError<'db>> { + let (enter_method, exit_method) = match mode { + EvaluationMode::Async => ("__aenter__", "__aexit__"), + EvaluationMode::Sync => ("__enter__", "__exit__"), + }; + + let enter = self.try_call_dunder(db, enter_method, CallArguments::none()); let exit = self.try_call_dunder( db, - "__exit__", + exit_method, CallArguments::positional([Type::none(db), Type::none(db), Type::none(db)]), ); // TODO: Make use of Protocols when we support it (the manager be assignable to `contextlib.AbstractContextManager`). match (enter, exit) { - (Ok(enter), Ok(_)) => Ok(enter.return_type(db)), - (Ok(enter), Err(exit_error)) => Err(ContextManagerError::Exit { - enter_return_type: enter.return_type(db), - exit_error, - }), + (Ok(enter), Ok(_)) => { + let ty = enter.return_type(db); + Ok(if mode.is_async() { + ty.resolve_await(db) + } else { + ty + }) + } + (Ok(enter), Err(exit_error)) => { + let ty = enter.return_type(db); + Err(ContextManagerError::Exit { + enter_return_type: if mode.is_async() { + ty.resolve_await(db) + } else { + ty + }, + exit_error, + mode, + }) + } // TODO: Use the `exit_ty` to determine if any raised exception is suppressed. - (Err(enter_error), Ok(_)) => Err(ContextManagerError::Enter(enter_error)), + (Err(enter_error), Ok(_)) => Err(ContextManagerError::Enter(enter_error, mode)), (Err(enter_error), Err(exit_error)) => Err(ContextManagerError::EnterAndExit { enter_error, exit_error, + mode, }), } } - /// Similar to [`Self::try_enter`], but for async context managers. - fn aenter(self, db: &'db dyn Db) -> Type<'db> { - // TODO: Add proper error handling and rename this method to `try_aenter`. - self.try_call_dunder(db, "__aenter__", CallArguments::none()) - .map_or(Type::unknown(), |result| { - result.return_type(db).resolve_await(db) - }) - } - /// Resolve the type of an `await …` expression where `self` is the type of the awaitable. fn resolve_await(self, db: &'db dyn Db) -> Type<'db> { // TODO: Add proper error handling and rename this method to `try_await`. @@ -6894,14 +6919,16 @@ impl<'db> TypeVarBoundOrConstraints<'db> { /// Error returned if a type is not (or may not be) a context manager. #[derive(Debug)] enum ContextManagerError<'db> { - Enter(CallDunderError<'db>), + Enter(CallDunderError<'db>, EvaluationMode), Exit { enter_return_type: Type<'db>, exit_error: CallDunderError<'db>, + mode: EvaluationMode, }, EnterAndExit { enter_error: CallDunderError<'db>, exit_error: CallDunderError<'db>, + mode: EvaluationMode, }, } @@ -6910,18 +6937,20 @@ impl<'db> ContextManagerError<'db> { self.enter_type(db).unwrap_or(Type::unknown()) } - /// Returns the `__enter__` return type if it is known, - /// or `None` if the type never has a callable `__enter__` attribute + /// Returns the `__enter__` or `__aenter__` return type if it is known, + /// or `None` if the type never has a callable `__enter__` or `__aenter__` attribute fn enter_type(&self, db: &'db dyn Db) -> Option> { match self { Self::Exit { enter_return_type, exit_error: _, + mode: _, } => Some(*enter_return_type), - Self::Enter(enter_error) + Self::Enter(enter_error, _) | Self::EnterAndExit { enter_error, exit_error: _, + mode: _, } => match enter_error { CallDunderError::PossiblyUnbound(call_outcome) => { Some(call_outcome.return_type(db)) @@ -6944,6 +6973,17 @@ impl<'db> ContextManagerError<'db> { return; }; + let mode = match self { + Self::Exit { mode, .. } | Self::Enter(_, mode) | Self::EnterAndExit { mode, .. } => { + *mode + } + }; + + let (enter_method, exit_method) = match mode { + EvaluationMode::Async => ("__aenter__", "__aexit__"), + EvaluationMode::Sync => ("__enter__", "__exit__"), + }; + let format_call_dunder_error = |call_dunder_error: &CallDunderError<'db>, name: &str| { match call_dunder_error { CallDunderError::MethodNotAvailable => format!("it does not implement `{name}`"), @@ -6987,38 +7027,52 @@ impl<'db> ContextManagerError<'db> { Self::Exit { enter_return_type: _, exit_error, - } => format_call_dunder_error(exit_error, "__exit__"), - Self::Enter(enter_error) => format_call_dunder_error(enter_error, "__enter__"), + mode: _, + } => format_call_dunder_error(exit_error, exit_method), + Self::Enter(enter_error, _) => format_call_dunder_error(enter_error, enter_method), Self::EnterAndExit { enter_error, exit_error, - } => format_call_dunder_errors(enter_error, "__enter__", exit_error, "__exit__"), + mode: _, + } => format_call_dunder_errors(enter_error, enter_method, exit_error, exit_method), }; - let mut diag = builder.into_diagnostic( - format_args!( - "Object of type `{context_expression}` cannot be used with `with` because {formatted_errors}", - context_expression = context_expression_type.display(db) - ), + // Suggest using `async with` if only async methods are available in a sync context, + // or suggest using `with` if only sync methods are available in an async context. + let with_kw = match mode { + EvaluationMode::Sync => "with", + EvaluationMode::Async => "async with", + }; + + let mut diag = builder.into_diagnostic(format_args!( + "Object of type `{}` cannot be used with `{}` because {}", + context_expression_type.display(db), + with_kw, + formatted_errors, + )); + + let (alt_mode, alt_enter_method, alt_exit_method, alt_with_kw) = match mode { + EvaluationMode::Sync => ("async", "__aenter__", "__aexit__", "async with"), + EvaluationMode::Async => ("sync", "__enter__", "__exit__", "with"), + }; + + let alt_enter = + context_expression_type.try_call_dunder(db, alt_enter_method, CallArguments::none()); + let alt_exit = context_expression_type.try_call_dunder( + db, + alt_exit_method, + CallArguments::positional([Type::unknown(), Type::unknown(), Type::unknown()]), ); - // If `__aenter__` and `__aexit__` are available, the user may have intended to use `async with` instead of `with`: - if let ( - Ok(_) | Err(CallDunderError::CallError(..)), - Ok(_) | Err(CallDunderError::CallError(..)), - ) = ( - context_expression_type.try_call_dunder(db, "__aenter__", CallArguments::none()), - context_expression_type.try_call_dunder( - db, - "__aexit__", - CallArguments::positional([Type::unknown(), Type::unknown(), Type::unknown()]), - ), - ) { + if (alt_enter.is_ok() || matches!(alt_enter, Err(CallDunderError::CallError(..)))) + && (alt_exit.is_ok() || matches!(alt_exit, Err(CallDunderError::CallError(..)))) + { diag.info(format_args!( - "Objects of type `{context_expression}` can be used as async context managers", - context_expression = context_expression_type.display(db) + "Objects of type `{}` can be used as {} context managers", + context_expression_type.display(db), + alt_mode )); - diag.info("Consider using `async with` here"); + diag.info(format!("Consider using `{alt_with_kw}` here")); } } } diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index dc9d9fe8ab..64464d0a1b 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -2583,7 +2583,11 @@ impl<'db> ClassLiteral<'db> { db, index.expression(with_item.context_expr(&module)), ); - let inferred_ty = context_ty.enter(db); + let inferred_ty = if with_item.is_async() { + context_ty.aenter(db) + } else { + context_ty.enter(db) + }; union_of_inferred_types = union_of_inferred_types.add(inferred_ty); } diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 60107e9c84..d6e6a91f59 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -3236,12 +3236,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { context_expression_type: Type<'db>, is_async: bool, ) -> Type<'db> { - if is_async { - return context_expression_type.aenter(self.db()); - } + let eval_mode = if is_async { + EvaluationMode::Async + } else { + EvaluationMode::Sync + }; context_expression_type - .try_enter(self.db()) + .try_enter_with_mode(self.db(), eval_mode) .unwrap_or_else(|err| { err.report_diagnostic( &self.context, diff --git a/crates/ty_python_semantic/src/types/unpacker.rs b/crates/ty_python_semantic/src/types/unpacker.rs index f1ebc42967..65bec0ef4f 100644 --- a/crates/ty_python_semantic/src/types/unpacker.rs +++ b/crates/ty_python_semantic/src/types/unpacker.rs @@ -75,20 +75,16 @@ impl<'db, 'ast> Unpacker<'db, 'ast> { ); err.fallback_element_type(self.db()) }), - UnpackKind::ContextManager { mode } => { - if mode.is_async() { - value_type.aenter(self.db()) - } else { - value_type.try_enter(self.db()).unwrap_or_else(|err| { - err.report_diagnostic( - &self.context, - value_type, - value.as_any_node_ref(self.db(), self.module()), - ); - err.fallback_enter_type(self.db()) - }) - } - } + UnpackKind::ContextManager { mode } => value_type + .try_enter_with_mode(self.db(), mode) + .unwrap_or_else(|err| { + err.report_diagnostic( + &self.context, + value_type, + value.as_any_node_ref(self.db(), self.module()), + ); + err.fallback_enter_type(self.db()) + }), }; self.unpack_inner(