mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-09-29 13:24:57 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[ty] Diagnostics for async context managers (#19704)

Browse source 
## Summary

Implements diagnostics for async context managers. Fixes
https://github.com/astral-sh/ty/issues/918.

## Test Plan

Mdtests have been added.
This commit is contained in:
Simon Lamon 2025-08-05 16:41:37 +02:00 • committed by GitHub
parent 78e5fe0a51
commit 934fd37d2b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 341 additions and 121 deletions
Download patch file Download diff file Expand all files Collapse all files

41
crates/ty_python_semantic/resources/mdtest/snapshots/async.md_-_Async_with_statement…_-_Accidental_use_of_as…_(5b8c1b4d846bc544).snap Normal file
View file

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

60
crates/ty_python_semantic/resources/mdtest/snapshots/sync.md_-_With_statements_-_Accidental_use_of_no…_(b07503f9b773ea61).snap
View file

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

165
crates/ty_python_semantic/resources/mdtest/with/async.md
View file

@ -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

6
crates/ty_python_semantic/resources/mdtest/with/sync.md
View file

@ -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

150
crates/ty_python_semantic/src/types.rs
View file

@ -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>> {
let enter = self.try_call_dunder(db, "__enter__", CallArguments::none());
fn try_enter_with_mode(
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), 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<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__"),
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"));
}
}
}

6
crates/ty_python_semantic/src/types/class.rs
View file

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

10
crates/ty_python_semantic/src/types/infer.rs
View file

@ -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,

24
crates/ty_python_semantic/src/types/unpacker.rs
View file

@ -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(