From 5bc0d9c3243458ce6eafb26da58571f7aa72ab07 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 12 Feb 2024 20:09:39 -0500 Subject: [PATCH] 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. --- crates/ruff_linter/src/checkers/ast/mod.rs | 53 ++++++++++++++++--- crates/ruff_linter/src/renamer.rs | 1 + .../src/rules/pandas_vet/helpers.rs | 1 + .../src/rules/pylint/rules/non_ascii_name.rs | 3 ++ .../src/analyze/typing.rs | 34 +++--------- crates/ruff_python_semantic/src/binding.rs | 6 +++ crates/ruff_python_semantic/src/model.rs | 28 ++++++++++ 7 files changed, 91 insertions(+), 35 deletions(-) diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs index 3d74001ecf..ce3083f2b1 100644 --- a/crates/ruff_linter/src/checkers/ast/mod.rs +++ b/crates/ruff_linter/src/checkers/ast/mod.rs @@ -1285,6 +1285,16 @@ where self.semantic.flags |= SemanticModelFlags::F_STRING; visitor::walk_expr(self, expr); } + Expr::NamedExpr(ast::ExprNamedExpr { + target, + value, + range: _, + }) => { + self.visit_expr(value); + + self.semantic.flags |= SemanticModelFlags::NAMED_EXPRESSION_ASSIGNMENT; + self.visit_expr(target); + } _ => visitor::walk_expr(self, expr), } @@ -1501,6 +1511,8 @@ impl<'a> Checker<'a> { unreachable!("Generator expression must contain at least one generator"); }; + let flags = self.semantic.flags; + // Generators are compiled as nested functions. (This may change with PEP 709.) // As such, the `iter` of the first generator is evaluated in the outer scope, while all // subsequent nodes are evaluated in the inner scope. @@ -1530,14 +1542,22 @@ impl<'a> Checker<'a> { // `x` is local to `foo`, and the `T` in `y=T` skips the class scope when resolving. self.visit_expr(&generator.iter); self.semantic.push_scope(ScopeKind::Generator); + + self.semantic.flags = flags | SemanticModelFlags::COMPREHENSION_ASSIGNMENT; self.visit_expr(&generator.target); + self.semantic.flags = flags; + for expr in &generator.ifs { self.visit_boolean_test(expr); } for generator in iterator { self.visit_expr(&generator.iter); + + self.semantic.flags = flags | SemanticModelFlags::COMPREHENSION_ASSIGNMENT; self.visit_expr(&generator.target); + self.semantic.flags = flags; + for expr in &generator.ifs { self.visit_boolean_test(expr); } @@ -1736,11 +1756,21 @@ impl<'a> Checker<'a> { return; } + // A binding within a `for` must be a loop variable, as in: + // ```python + // for x in range(10): + // ... + // ``` if parent.is_for_stmt() { self.add_binding(id, expr.range(), BindingKind::LoopVar, flags); return; } + // A binding within a `with` must be an item, as in: + // ```python + // with open("file.txt") as fp: + // ... + // ``` if parent.is_with_stmt() { self.add_binding(id, expr.range(), BindingKind::WithItemVar, flags); return; @@ -1796,17 +1826,26 @@ impl<'a> Checker<'a> { } // If the expression is the left-hand side of a walrus operator, then it's a named - // expression assignment. - if self - .semantic - .current_expressions() - .filter_map(Expr::as_named_expr_expr) - .any(|parent| parent.target.as_ref() == expr) - { + // expression assignment, as in: + // ```python + // if (x := 10) > 5: + // ... + // ``` + if self.semantic.in_named_expression_assignment() { self.add_binding(id, expr.range(), BindingKind::NamedExprAssignment, flags); return; } + // If the expression is part of a comprehension target, then it's a comprehension variable + // assignment, as in: + // ```python + // [x for x in range(10)] + // ``` + if self.semantic.in_comprehension_assignment() { + self.add_binding(id, expr.range(), BindingKind::ComprehensionVar, flags); + return; + } + self.add_binding(id, expr.range(), BindingKind::Assignment, flags); } diff --git a/crates/ruff_linter/src/renamer.rs b/crates/ruff_linter/src/renamer.rs index fd2cea8304..56fb1e6ad0 100644 --- a/crates/ruff_linter/src/renamer.rs +++ b/crates/ruff_linter/src/renamer.rs @@ -248,6 +248,7 @@ impl Renamer { | BindingKind::Assignment | BindingKind::BoundException | BindingKind::LoopVar + | BindingKind::ComprehensionVar | BindingKind::WithItemVar | BindingKind::Global | BindingKind::Nonlocal(_) diff --git a/crates/ruff_linter/src/rules/pandas_vet/helpers.rs b/crates/ruff_linter/src/rules/pandas_vet/helpers.rs index 85ba34ec77..c88555659d 100644 --- a/crates/ruff_linter/src/rules/pandas_vet/helpers.rs +++ b/crates/ruff_linter/src/rules/pandas_vet/helpers.rs @@ -47,6 +47,7 @@ pub(super) fn test_expression(expr: &Expr, semantic: &SemanticModel) -> Resoluti | BindingKind::Assignment | BindingKind::NamedExprAssignment | BindingKind::LoopVar + | BindingKind::ComprehensionVar | BindingKind::Global | BindingKind::Nonlocal(_) => Resolution::RelevantLocal, BindingKind::Import(import) if matches!(import.call_path(), ["pandas"]) => { diff --git a/crates/ruff_linter/src/rules/pylint/rules/non_ascii_name.rs b/crates/ruff_linter/src/rules/pylint/rules/non_ascii_name.rs index 93733827b4..92bcd79907 100644 --- a/crates/ruff_linter/src/rules/pylint/rules/non_ascii_name.rs +++ b/crates/ruff_linter/src/rules/pylint/rules/non_ascii_name.rs @@ -52,6 +52,7 @@ pub(crate) fn non_ascii_name(binding: &Binding, locator: &Locator) -> Option Kind::Assignment, BindingKind::TypeParam => Kind::TypeParam, BindingKind::LoopVar => Kind::LoopVar, + BindingKind::ComprehensionVar => Kind::ComprenhensionVar, BindingKind::WithItemVar => Kind::WithItemVar, BindingKind::Global => Kind::Global, BindingKind::Nonlocal(_) => Kind::Nonlocal, @@ -88,6 +89,7 @@ enum Kind { Assignment, TypeParam, LoopVar, + ComprenhensionVar, WithItemVar, Global, Nonlocal, @@ -105,6 +107,7 @@ impl fmt::Display for Kind { Kind::Assignment => f.write_str("Variable"), Kind::TypeParam => f.write_str("Type parameter"), Kind::LoopVar => f.write_str("Variable"), + Kind::ComprenhensionVar => f.write_str("Variable"), Kind::WithItemVar => f.write_str("Variable"), Kind::Global => f.write_str("Global"), Kind::Nonlocal => f.write_str("Nonlocal"), diff --git a/crates/ruff_python_semantic/src/analyze/typing.rs b/crates/ruff_python_semantic/src/analyze/typing.rs index cb1a1c15eb..3283db129d 100644 --- a/crates/ruff_python_semantic/src/analyze/typing.rs +++ b/crates/ruff_python_semantic/src/analyze/typing.rs @@ -426,16 +426,8 @@ fn check_type(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(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(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, }, diff --git a/crates/ruff_python_semantic/src/binding.rs b/crates/ruff_python_semantic/src/binding.rs index e2f29ae07f..9057c60de1 100644 --- a/crates/ruff_python_semantic/src/binding.rs +++ b/crates/ruff_python_semantic/src/binding.rs @@ -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: diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs index 516c59d1f6..add977d754 100644 --- a/crates/ruff_python_semantic/src/model.rs +++ b/crates/ruff_python_semantic/src/model.rs @@ -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();