Add a binding kind for comprehension targets (#9967)

## Summary

I was surprised to learn that we treat `x` in `[_ for x in y]` as an
"assignment" binding kind, rather than a dedicated comprehension
variable.
This commit is contained in:
Charlie Marsh 2024-02-12 20:09:39 -05:00 committed by GitHub
parent cf77eeb913
commit 5bc0d9c324
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 91 additions and 35 deletions

View file

@ -426,16 +426,8 @@ fn check_type<T: TypeChecker>(binding: &Binding, semantic: &SemanticModel) -> bo
// ```
//
// The type checker might know how to infer the type based on `init_expr`.
Some(Stmt::Assign(ast::StmtAssign { targets, value, .. })) => {
// TODO(charlie): Replace this with `find_binding_value`, which matches the values.
if targets
.iter()
.any(|target| target.range().contains_range(binding.range()))
{
T::match_initializer(value.as_ref(), semantic)
} else {
false
}
Some(Stmt::Assign(ast::StmtAssign { value, .. })) => {
T::match_initializer(value.as_ref(), semantic)
}
// ```python
@ -443,15 +435,8 @@ fn check_type<T: TypeChecker>(binding: &Binding, semantic: &SemanticModel) -> bo
// ```
//
// In this situation, we check only the annotation.
Some(Stmt::AnnAssign(ast::StmtAnnAssign {
target, annotation, ..
})) => {
// TODO(charlie): Replace this with `find_binding_value`, which matches the values.
if target.range().contains_range(binding.range()) {
T::match_annotation(annotation.as_ref(), semantic)
} else {
false
}
Some(Stmt::AnnAssign(ast::StmtAnnAssign { annotation, .. })) => {
T::match_annotation(annotation.as_ref(), semantic)
}
_ => false,
},
@ -481,15 +466,8 @@ fn check_type<T: TypeChecker>(binding: &Binding, semantic: &SemanticModel) -> bo
// ```
//
// It's a typed declaration, type annotation is the only source of information.
Some(Stmt::AnnAssign(ast::StmtAnnAssign {
target, annotation, ..
})) => {
// TODO(charlie): Replace this with `find_binding_value`, which matches the values.
if target.range().contains_range(binding.range()) {
T::match_annotation(annotation.as_ref(), semantic)
} else {
false
}
Some(Stmt::AnnAssign(ast::StmtAnnAssign { annotation, .. })) => {
T::match_annotation(annotation.as_ref(), semantic)
}
_ => false,
},

View file

@ -432,6 +432,12 @@ pub enum BindingKind<'a> {
/// ```
LoopVar,
/// A binding for a comprehension variable, like `x` in:
/// ```python
/// [x for x in range(10)]
/// ```
ComprehensionVar,
/// A binding for a with statement variable, like `x` in:
/// ```python
/// with open('foo.py') as x:

View file

@ -1511,6 +1511,18 @@ impl<'a> SemanticModel<'a> {
.intersects(SemanticModelFlags::FUTURE_ANNOTATIONS)
}
/// Return `true` if the model is in a named expression assignment (e.g., `x := 1`).
pub const fn in_named_expression_assignment(&self) -> bool {
self.flags
.intersects(SemanticModelFlags::NAMED_EXPRESSION_ASSIGNMENT)
}
/// Return `true` if the model is in a comprehension assignment (e.g., `_ for x in y`).
pub const fn in_comprehension_assignment(&self) -> bool {
self.flags
.intersects(SemanticModelFlags::COMPREHENSION_ASSIGNMENT)
}
/// Return an iterator over all bindings shadowed by the given [`BindingId`], within the
/// containing scope, and across scopes.
pub fn shadowed_bindings(
@ -1825,6 +1837,22 @@ bitflags! {
///
const TYPE_PARAM_DEFINITION = 1 << 17;
/// The model is in a named expression assignment.
///
/// For example, the model could be visiting `x` in:
/// ```python
/// if (x := 1): ...
/// ```
const NAMED_EXPRESSION_ASSIGNMENT = 1 << 18;
/// The model is in a comprehension variable assignment.
///
/// For example, the model could be visiting `x` in:
/// ```python
/// [_ for x in range(10)]
/// ```
const COMPREHENSION_ASSIGNMENT = 1 << 19;
/// The context is in any type annotation.
const ANNOTATION = Self::TYPING_ONLY_ANNOTATION.bits() | Self::RUNTIME_EVALUATED_ANNOTATION.bits() | Self::RUNTIME_REQUIRED_ANNOTATION.bits();