Fix scoping of comprehensions within classes (#4494)

This commit is contained in:
Charlie Marsh 2023-05-18 10:30:02 -04:00 committed by GitHub
parent e8e66f3824
commit 2fb312bb2b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 134 additions and 29 deletions

View file

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

View file

@ -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<String> {
let content = locator.slice(expr.range());
let mut expression = match_expression(content)?;

View file

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

View file

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