Improve fix, use Checker.current_scopes()

exit inside functions is still not detected

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)

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<Item = &Scope> {
+        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!(

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<String> {
-    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<String> {
+    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,

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