[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.

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
+
+```

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:
+1 | class Manager:
- 2 |     async def __aenter__(self): ...
+2 |     async def __aenter__(self): ...
- 3 |     async def __aexit__(self, *args): ...
+3 |     async def __aexit__(self, *args): ...
- 4 | 
+4 | 
- 5 | # error: [invalid-context-manager] "Object of type `Manager` cannot be used with `with` because it does not implement `__enter__` and `__exit__`"
+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():
+6 | with Manager():
- 7 |     ...
+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 |     ...
 ```
 
 # 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
-
-```

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`
+
+<!-- snapshot-diagnostics -->
+
+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

crates/ty_python_semantic/resources/mdtest/with/sync.md

@@ -155,7 +155,7 @@ with context_expr as f:
 <!-- snapshot-diagnostics -->
 
 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

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<Type<'db>, ContextManagerError<'db>> {
+    fn try_enter_with_mode(
-        let enter = self.try_call_dunder(db, "__enter__", CallArguments::none());
+        self,
+        db: &'db dyn Db,
+        mode: EvaluationMode,
+    ) -> Result<Type<'db>, 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), Ok(_)) => {
-            (Ok(enter), Err(exit_error)) => Err(ContextManagerError::Exit {
+                let ty = enter.return_type(db);
-                enter_return_type: enter.return_type(db),
+                Ok(if mode.is_async() {
-                exit_error,
+                    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,
+    /// Returns the `__enter__` or `__aenter__` return type if it is known,
-    /// or `None` if the type never has a callable `__enter__` attribute
+    /// or `None` if the type never has a callable `__enter__` or `__aenter__` attribute
     fn enter_type(&self, db: &'db dyn Db) -> Option<Type<'db>> {
         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__"),
+                mode: _,
-            Self::Enter(enter_error) => format_call_dunder_error(enter_error, "__enter__"),
+            } => 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(
+        // Suggest using `async with` if only async methods are available in a sync context,
-            format_args!(
+        // or suggest using `with` if only sync methods are available in an async context.
-                "Object of type `{context_expression}` cannot be used with `with` because {formatted_errors}",
+        let with_kw = match mode {
-                context_expression = context_expression_type.display(db)
+            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 (alt_enter.is_ok() || matches!(alt_enter, Err(CallDunderError::CallError(..))))
-        if let (
+            && (alt_exit.is_ok() || matches!(alt_exit, Err(CallDunderError::CallError(..))))
-            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()]),
-            ),
-        ) {
             diag.info(format_args!(
-                "Objects of type `{context_expression}` can be used as async context managers",
+                "Objects of type `{}` can be used as {} context managers",
-                context_expression = context_expression_type.display(db)
+                context_expression_type.display(db),
+                alt_mode
             ));
-            diag.info("Consider using `async with` here");
+            diag.info(format!("Consider using `{alt_with_kw}` here"));
         }
     }
 }

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);
                             }

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 {
+        let eval_mode = if is_async {
-            return context_expression_type.aenter(self.db());
+            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,

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 } => {
+            UnpackKind::ContextManager { mode } => value_type
-                if mode.is_async() {
+                .try_enter_with_mode(self.db(), mode)
-                    value_type.aenter(self.db())
+                .unwrap_or_else(|err| {
-                } else {
+                    err.report_diagnostic(
-                    value_type.try_enter(self.db()).unwrap_or_else(|err| {
+                        &self.context,
-                        err.report_diagnostic(
+                        value_type,
-                            &self.context,
+                        value.as_any_node_ref(self.db(), self.module()),
-                            value_type,
+                    );
-                            value.as_any_node_ref(self.db(), self.module()),
+                    err.fallback_enter_type(self.db())
-                        );
+                }),
-                        err.fallback_enter_type(self.db())
-                    })
-                }
-            }
         };
 
         self.unpack_inner(