From 2fb312bb2befec66e93a66e70a37e6f91d27a7cb Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 18 May 2023 10:30:02 -0400 Subject: [PATCH] Fix scoping of comprehensions within classes (#4494) --- crates/ruff/src/checkers/ast/mod.rs | 135 +++++++++++++++--- .../flake8_simplify/rules/key_in_dict.rs | 2 +- crates/ruff/src/rules/pyflakes/mod.rs | 17 +++ .../rules/ruff/rules/pairwise_over_zipped.rs | 9 +- 4 files changed, 134 insertions(+), 29 deletions(-) diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index d5e767b8442a3..44c50698f3676 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/crates/ruff/src/checkers/ast/mod.rs @@ -3731,7 +3731,15 @@ where if self.settings.rules.enabled(Rule::FunctionUsesLoopVariable) { flake8_bugbear::rules::function_uses_loop_variable(self, &Node::Expr(expr)); } - self.ctx.push_scope(ScopeKind::Generator); + if self.settings.rules.enabled(Rule::InDictKeys) { + for generator in generators { + flake8_simplify::rules::key_in_dict_for( + self, + &generator.target, + &generator.iter, + ); + } + } } Expr::DictComp(ast::ExprDictComp { key, @@ -3747,13 +3755,33 @@ where if self.settings.rules.enabled(Rule::FunctionUsesLoopVariable) { flake8_bugbear::rules::function_uses_loop_variable(self, &Node::Expr(expr)); } - self.ctx.push_scope(ScopeKind::Generator); + if self.settings.rules.enabled(Rule::InDictKeys) { + for generator in generators { + flake8_simplify::rules::key_in_dict_for( + self, + &generator.target, + &generator.iter, + ); + } + } } - Expr::GeneratorExp(_) => { + Expr::GeneratorExp(ast::ExprGeneratorExp { + generators, + elt: _, + range: _, + }) => { if self.settings.rules.enabled(Rule::FunctionUsesLoopVariable) { flake8_bugbear::rules::function_uses_loop_variable(self, &Node::Expr(expr)); } - self.ctx.push_scope(ScopeKind::Generator); + if self.settings.rules.enabled(Rule::InDictKeys) { + for generator in generators { + flake8_simplify::rules::key_in_dict_for( + self, + &generator.target, + &generator.iter, + ); + } + } } Expr::BoolOp(ast::ExprBoolOp { op, @@ -3801,6 +3829,34 @@ where // Recurse. match expr { + Expr::ListComp(ast::ExprListComp { + elt, + generators, + range: _, + }) + | Expr::SetComp(ast::ExprSetComp { + elt, + generators, + range: _, + }) + | Expr::GeneratorExp(ast::ExprGeneratorExp { + elt, + generators, + range: _, + }) => { + self.visit_generators(generators); + self.visit_expr(elt); + } + Expr::DictComp(ast::ExprDictComp { + key, + value, + generators, + range: _, + }) => { + self.visit_generators(generators); + self.visit_expr(key); + self.visit_expr(value); + } Expr::Lambda(_) => { self.deferred.lambdas.push((expr, self.ctx.snapshot())); } @@ -4079,21 +4135,6 @@ where self.ctx.pop_expr(); } - fn visit_comprehension(&mut self, comprehension: &'b Comprehension) { - if self.settings.rules.enabled(Rule::InDictKeys) { - flake8_simplify::rules::key_in_dict_for( - self, - &comprehension.target, - &comprehension.iter, - ); - } - self.visit_expr(&comprehension.iter); - self.visit_expr(&comprehension.target); - for expr in &comprehension.ifs { - self.visit_boolean_test(expr); - } - } - fn visit_excepthandler(&mut self, excepthandler: &'b Excepthandler) { match excepthandler { Excepthandler::ExceptHandler(ast::ExcepthandlerExceptHandler { @@ -4398,6 +4439,58 @@ impl<'a> Checker<'a> { docstring.is_some() } + /// Visit a list of [`Comprehension`] nodes, assumed to be the comprehensions that compose a + /// generator expression, like a list or set comprehension. + fn visit_generators(&mut self, generators: &'a [Comprehension]) { + let mut generators = generators.iter(); + + let Some(generator) = generators.next() else { + unreachable!("Generator expression must contain at least one generator"); + }; + + // 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. + // + // For example, given: + // ```py + // class A: + // T = range(10) + // + // L = [x for x in T for y in T] + // ``` + // + // Conceptually, this is compiled as: + // ```py + // class A: + // T = range(10) + // + // def foo(x=T): + // def bar(y=T): + // pass + // return bar() + // foo() + // ``` + // + // Following Python's scoping rules, the `T` in `x=T` is thus evaluated in the outer scope, + // while all subsequent reads and writes are evaluated in the inner scope. In particular, + // `x` is local to `foo`, and the `T` in `y=T` skips the class scope when resolving. + self.visit_expr(&generator.iter); + self.ctx.push_scope(ScopeKind::Generator); + self.visit_expr(&generator.target); + for expr in &generator.ifs { + self.visit_boolean_test(expr); + } + + for generator in generators { + self.visit_expr(&generator.iter); + self.visit_expr(&generator.target); + for expr in &generator.ifs { + self.visit_boolean_test(expr); + } + } + } + /// Visit an body of [`Stmt`] nodes within a type-checking block. fn visit_type_checking_block(&mut self, body: &'a [Stmt]) { let snapshot = self.ctx.flags; @@ -4614,14 +4707,13 @@ impl<'a> Checker<'a> { let id = id.as_str(); let mut first_iter = true; - let mut in_generator = false; let mut import_starred = false; for scope in self.ctx.scopes.ancestors(self.ctx.scope_id) { if scope.kind.is_class() { if id == "__class__" { return; - } else if !first_iter && !in_generator { + } else if !first_iter { continue; } } @@ -4689,7 +4781,6 @@ impl<'a> Checker<'a> { } first_iter = false; - in_generator = matches!(scope.kind, ScopeKind::Generator); import_starred = import_starred || scope.uses_star_imports(); } diff --git a/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs b/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs index 7158040d67506..d986ecd9b4fb2 100644 --- a/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs +++ b/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs @@ -35,7 +35,7 @@ impl AlwaysAutofixableViolation for InDictKeys { fn get_value_content_for_key_in_dict( locator: &Locator, stylist: &Stylist, - expr: &rustpython_parser::ast::Expr, + expr: &Expr, ) -> Result { let content = locator.slice(expr.range()); let mut expression = match_expression(content)?; diff --git a/crates/ruff/src/rules/pyflakes/mod.rs b/crates/ruff/src/rules/pyflakes/mod.rs index dcdbedebb0f92..648f93ba118e0 100644 --- a/crates/ruff/src/rules/pyflakes/mod.rs +++ b/crates/ruff/src/rules/pyflakes/mod.rs @@ -1030,6 +1030,23 @@ mod tests { "#, &[], ); + + flakes( + r#" + class A: + T = 1 + + # In each case, `T` is undefined. Only the first `iter` uses the class scope. + X = (T for x in range(10)) + Y = [x for x in range(10) if T] + Z = [x for x in range(10) for y in T] + "#, + &[ + Rule::UndefinedName, + Rule::UndefinedName, + Rule::UndefinedName, + ], + ); } #[test] diff --git a/crates/ruff/src/rules/ruff/rules/pairwise_over_zipped.rs b/crates/ruff/src/rules/ruff/rules/pairwise_over_zipped.rs index d7f7219d58af7..e93a5d5a2926f 100644 --- a/crates/ruff/src/rules/ruff/rules/pairwise_over_zipped.rs +++ b/crates/ruff/src/rules/ruff/rules/pairwise_over_zipped.rs @@ -99,10 +99,7 @@ pub(crate) fn pairwise_over_zipped(checker: &mut Checker, func: &Expr, args: &[E } // Require the function to be the builtin `zip`. - if id != "zip" { - return; - } - if !checker.ctx.is_builtin(id) { + if !(id == "zip" && checker.ctx.is_builtin(id)) { return; } @@ -118,9 +115,9 @@ pub(crate) fn pairwise_over_zipped(checker: &mut Checker, func: &Expr, args: &[E }; // Require second argument to be a `Subscript`. - let Expr::Subscript (_) = &args[1] else { + if !matches!(&args[1], Expr::Subscript(_)) { return; - }; + } let Some(second_arg_info) = match_slice_info(&args[1]) else { return; };