diff --git a/resources/test/fixtures/RUF101.py b/resources/test/fixtures/RUF101.py index faf0d907af5b66..669f3934681291 100644 --- a/resources/test/fixtures/RUF101.py +++ b/resources/test/fixtures/RUF101.py @@ -1,6 +1,10 @@ exit(1) +def main(): + exit(6) # Is not detected + + import sys exit(0) @@ -10,7 +14,7 @@ def main(): exit(3) # Is not detected -sys.exit(0) +sys.exit(7) del sys @@ -19,8 +23,24 @@ def main(): exit(2) +def main(): + exit(9) # Is not detected + + +del sys2 + + +from sys import exit as exit2 + +exit(5) + + +def main(): + exit(10) # Is not detected + + def exit(e): pass -exit(0) +exit(8) diff --git a/src/check_ast.rs b/src/check_ast.rs index cdadbe344f464d..6a40b263c604de 100644 --- a/src/check_ast.rs +++ b/src/check_ast.rs @@ -418,7 +418,7 @@ where } StmtKind::Return { .. } => { if self.settings.enabled.contains(&CheckCode::F706) { - if let Some(index) = self.scope_stack.last().cloned() { + if let Some(&index) = self.scope_stack.last() { if matches!( self.scopes[index].kind, ScopeKind::Class(_) | ScopeKind::Module @@ -2124,6 +2124,10 @@ impl<'a> Checker<'a> { &self.scopes[*(self.scope_stack.last().expect("No current scope found."))] } + pub fn current_scopes(&self) -> impl Iterator { + self.scope_stack.iter().rev().map(|s| &self.scopes[*s]) + } + pub fn current_parent(&self) -> &'a Stmt { self.parents[*(self.parent_stack.last().expect("No parent found."))] } @@ -2143,7 +2147,7 @@ impl<'a> Checker<'a> { 'b: 'a, { if self.settings.enabled.contains(&CheckCode::F402) { - let scope = &self.scopes[*(self.scope_stack.last().expect("No current scope found."))]; + let scope = self.current_scope(); if let Some(existing) = scope.values.get(&name) { if matches!(binding.kind, BindingKind::LoopVar) && matches!( @@ -2168,7 +2172,7 @@ impl<'a> Checker<'a> { // TODO(charlie): Don't treat annotations as assignments if there is an existing // value. - let scope = &self.scopes[*(self.scope_stack.last().expect("No current scope found."))]; + let scope = self.current_scope(); let binding = match scope.values.get(&name) { None => binding, Some(existing) => Binding { @@ -2184,8 +2188,7 @@ impl<'a> Checker<'a> { fn handle_node_load(&mut self, expr: &Expr) { if let ExprKind::Name { id, .. } = &expr.node { - let scope_id = - self.scopes[*(self.scope_stack.last().expect("No current scope found."))].id; + let scope_id = self.current_scope().id; let mut first_iter = true; let mut in_generator = false; @@ -2340,7 +2343,7 @@ impl<'a> Checker<'a> { return; } - let current = &self.scopes[*(self.scope_stack.last().expect("No current scope found."))]; + let current = self.current_scope(); if id == "__all__" && matches!(current.kind, ScopeKind::Module) && matches!( diff --git a/src/rules/plugins/convert_sys_to_sys_exit.rs b/src/rules/plugins/convert_sys_to_sys_exit.rs index 57c03b6caaa1da..c04e4f9cb12246 100644 --- a/src/rules/plugins/convert_sys_to_sys_exit.rs +++ b/src/rules/plugins/convert_sys_to_sys_exit.rs @@ -1,35 +1,44 @@ use rustpython_ast::{Expr, ExprKind}; -use crate::ast::types::{Binding, BindingKind, Range, Scope}; +use crate::ast::types::{Binding, BindingKind, Range}; use crate::autofix::Fix; use crate::check_ast::Checker; use crate::checks::{Check, CheckKind}; -fn get_sys_import(scope: &Scope) -> Option { - let sys_binding_kind = - scope.values.values().map(|v| &v.kind).find( - |b| matches!(b, BindingKind::Importation(_, sys, _) if sys == &"sys".to_string()), - ); - if let Some(BindingKind::Importation(name, ..)) = sys_binding_kind { - return Some(name.clone()); - } - None +fn has_builtin_exit_in_scope(checker: &Checker) -> bool { + matches!( + checker.current_scopes().find_map(|s| s.values.get("exit")), + Some(Binding { + kind: BindingKind::Builtin, + .. + }) + ) +} + +fn get_exit_name(checker: &Checker) -> Option { + checker.current_scopes().find_map(|s| { + s.values.values().find_map(|v| match &v.kind { + BindingKind::Importation(name, full_name, _) if full_name == &"sys".to_string() => { + Some(name.to_owned() + ".exit") + } + BindingKind::FromImportation(name, full_name, _) + if full_name == &"sys.exit".to_string() => + { + Some(name.to_owned()) + } + _ => None, + }) + }) } pub fn convert_sys_to_sys_exit(checker: &mut Checker, func: &Expr) { if let ExprKind::Name { id, .. } = &func.node { if id == "exit" { - let scope = checker.current_scope(); - if let Some(Binding { - kind: BindingKind::Builtin, - .. - }) = scope.values.get("exit") - { + if has_builtin_exit_in_scope(checker) { let mut check = Check::new(CheckKind::ConvertExitToSysExit, Range::from_located(func)); if checker.patch(check.kind.code()) { - if let Some(mut content) = get_sys_import(scope) { - content.push_str(".exit"); + if let Some(content) = get_exit_name(checker) { check.amend(Fix::replacement( content, func.location, diff --git a/src/snapshots/ruff__linter__tests__RUF101_RUF101.py.snap b/src/snapshots/ruff__linter__tests__RUF101_RUF101.py.snap index 89063ed540cdda..e0e2278a38f088 100644 --- a/src/snapshots/ruff__linter__tests__RUF101_RUF101.py.snap +++ b/src/snapshots/ruff__linter__tests__RUF101_RUF101.py.snap @@ -12,36 +12,53 @@ expression: checks fix: ~ - kind: ConvertExitToSysExit location: - row: 6 + row: 10 column: 0 end_location: - row: 6 + row: 10 column: 4 fix: patch: content: sys.exit location: - row: 6 + row: 10 column: 0 end_location: - row: 6 + row: 10 column: 4 applied: false - kind: ConvertExitToSysExit location: - row: 19 + row: 23 column: 0 end_location: - row: 19 + row: 23 column: 4 fix: patch: content: sys2.exit location: - row: 19 + row: 23 column: 0 end_location: - row: 19 + row: 23 + column: 4 + applied: false +- kind: ConvertExitToSysExit + location: + row: 35 + column: 0 + end_location: + row: 35 + column: 4 + fix: + patch: + content: exit2 + location: + row: 35 + column: 0 + end_location: + row: 35 column: 4 applied: false