Correctly handle references in __all__ definitions when renaming sy…

…mbols in autofixes

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI025_2.py

@@ -1,3 +1,11 @@
+"""Tests to ensure we correctly rename references inside `__all__`"""
+
 from collections.abc import Set
 
 __all__ = ["Set"]
+
+if True:
+    __all__ += [r'''Set''']
+
+if 1:
+    __all__ += ["S" "e" "t"]

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI025_2.pyi

@@ -1,3 +1,11 @@
+"""Tests to ensure we correctly rename references inside `__all__`"""
+
 from collections.abc import Set
 
 __all__ = ["Set"]
+
+if True:
+    __all__ += [r'''Set''']
+
+if 1:
+    __all__ += ["S" "e" "t"]

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -2110,6 +2110,8 @@ impl<'a> Checker<'a> {
             .flatten()
             .collect();
 
+        self.semantic.flags |= SemanticModelFlags::DUNDER_ALL_DEFINITION;
+
         for DunderAllName { name, range } in exports {
             if let Some(binding_id) = self.semantic.global_scope().get(name) {
                 // Mark anything referenced in `__all__` as used.

crates/ruff_linter/src/renamer.rs

@@ -4,7 +4,9 @@ use anyhow::{anyhow, Result};
 use itertools::Itertools;
 
 use ruff_diagnostics::Edit;
+use ruff_python_ast::str::raw_contents_range;
 use ruff_python_semantic::{Binding, BindingKind, Scope, ScopeId, SemanticModel};
+use ruff_source_file::Locator;
 use ruff_text_size::Ranged;
 
 pub(crate) struct Renamer;
@@ -104,6 +106,7 @@ impl Renamer {
         target: &str,
         scope: &Scope,
         semantic: &SemanticModel,
+        locator: &Locator,
     ) -> Result<(Edit, Vec<Edit>)> {
         let mut edits = vec![];
 
@@ -130,7 +133,9 @@ impl Renamer {
         });
 
         let scope = scope_id.map_or(scope, |scope_id| &semantic.scopes[scope_id]);
-        edits.extend(Renamer::rename_in_scope(name, target, scope, semantic));
+        edits.extend(Renamer::rename_in_scope(
+            name, target, scope, semantic, locator,
+        ));
 
         // Find any scopes in which the symbol is referenced as `nonlocal` or `global`. For example,
         // given:
@@ -160,7 +165,9 @@ impl Renamer {
             .copied()
         {
             let scope = &semantic.scopes[scope_id];
-            edits.extend(Renamer::rename_in_scope(name, target, scope, semantic));
+            edits.extend(Renamer::rename_in_scope(
+                name, target, scope, semantic, locator,
+            ));
         }
 
         // Deduplicate any edits.
@@ -180,6 +187,7 @@ impl Renamer {
         target: &str,
         scope: &Scope,
         semantic: &SemanticModel,
+        locator: &Locator,
     ) -> Vec<Edit> {
         let mut edits = vec![];
 
@@ -202,7 +210,21 @@ impl Renamer {
                 // Rename the references to the binding.
                 edits.extend(binding.references().map(|reference_id| {
                     let reference = semantic.reference(reference_id);
-                    Edit::range_replacement(target.to_string(), reference.range())
+                    let range_to_replace = {
+                        if reference.in_dunder_all_definition() {
+                            let reference_source = locator.slice(reference.range());
+                            let relative_range = raw_contents_range(reference_source)
+                                .expect(
+                                    "Expected all references on the r.h.s. of an `__all__` definition to be strings"
+                                );
+                            debug_assert!(!relative_range.is_empty());
+                            relative_range + reference.start()
+                        } else {
+                            debug_assert!(locator.slice(reference.range()).contains(name));
+                            reference.range()
+                        }
+                    };
+                    Edit::range_replacement(target.to_string(), range_to_replace)
                 }));
             }
         }

crates/ruff_linter/src/rules/flake8_import_conventions/rules/unconventional_import_alias.rs

@@ -64,8 +64,9 @@ pub(crate) fn unconventional_import_alias(
     let import = binding.as_any_import()?;
     let qualified_name = import.qualified_name().to_string();
     let expected_alias = conventions.get(qualified_name.as_str())?;
+    let locator = checker.locator();
 
-    let name = binding.name(checker.locator());
+    let name = binding.name(locator);
     if binding.is_alias() && name == expected_alias {
         return None;
     }
@@ -82,7 +83,7 @@ pub(crate) fn unconventional_import_alias(
             diagnostic.try_set_fix(|| {
                 let scope = &checker.semantic().scopes[binding.scope];
                 let (edit, rest) =
-                    Renamer::rename(name, expected_alias, scope, checker.semantic())?;
+                    Renamer::rename(name, expected_alias, scope, checker.semantic(), locator)?;
                 Ok(Fix::unsafe_edits(edit, rest))
             });
         }

crates/ruff_linter/src/rules/flake8_pyi/rules/unaliased_collections_abc_set_import.rs

@@ -68,7 +68,8 @@ pub(crate) fn unaliased_collections_abc_set_import(
         return None;
     }
 
-    let name = binding.name(checker.locator());
+    let locator = checker.locator();
+    let name = binding.name(locator);
     if name == "AbstractSet" {
         return None;
     }
@@ -77,7 +78,8 @@ pub(crate) fn unaliased_collections_abc_set_import(
     if checker.semantic().is_available("AbstractSet") {
         diagnostic.try_set_fix(|| {
             let scope = &checker.semantic().scopes[binding.scope];
-            let (edit, rest) = Renamer::rename(name, "AbstractSet", scope, checker.semantic())?;
+            let (edit, rest) =
+                Renamer::rename(name, "AbstractSet", scope, checker.semantic(), locator)?;
             Ok(Fix::applicable_edits(
                 edit,
                 rest,

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI025_PYI025_2.py.snap

@@ -0,0 +1,30 @@
+---
+source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
+---
+PYI025_2.py:3:29: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+  |
+1 | """Tests to ensure we correctly rename references inside `__all__`"""
+2 | 
+3 | from collections.abc import Set
+  |                             ^^^ PYI025
+4 | 
+5 | __all__ = ["Set"]
+  |
+  = help: Alias `Set` to `AbstractSet`
+
+ℹ Unsafe fix
+1  1  | """Tests to ensure we correctly rename references inside `__all__`"""
+2  2  | 
+3     |-from collections.abc import Set
+   3  |+from collections.abc import Set as AbstractSet
+4  4  | 
+5     |-__all__ = ["Set"]
+   5  |+__all__ = ["AbstractSet"]
+6  6  | 
+7  7  | if True:
+8     |-    __all__ += [r'''Set''']
+   8  |+    __all__ += [r'''AbstractSet''']
+9  9  | 
+10 10 | if 1:
+11    |-    __all__ += ["S" "e" "t"]
+   11 |+    __all__ += ["AbstractSet"]

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI025_PYI025_2.pyi.snap

@@ -0,0 +1,30 @@
+---
+source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
+---
+PYI025_2.pyi:3:29: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+  |
+1 | """Tests to ensure we correctly rename references inside `__all__`"""
+2 | 
+3 | from collections.abc import Set
+  |                             ^^^ PYI025
+4 | 
+5 | __all__ = ["Set"]
+  |
+  = help: Alias `Set` to `AbstractSet`
+
+ℹ Unsafe fix
+1  1  | """Tests to ensure we correctly rename references inside `__all__`"""
+2  2  | 
+3     |-from collections.abc import Set
+   3  |+from collections.abc import Set as AbstractSet
+4  4  | 
+5     |-__all__ = ["Set"]
+   5  |+__all__ = ["AbstractSet"]
+6  6  | 
+7  7  | if True:
+8     |-    __all__ += [r'''Set''']
+   8  |+    __all__ += [r'''AbstractSet''']
+9  9  | 
+10 10 | if 1:
+11    |-    __all__ += ["S" "e" "t"]
+   11 |+    __all__ += ["AbstractSet"]

crates/ruff_linter/src/rules/pep8_naming/rules/invalid_first_argument_name.rs

@@ -6,6 +6,7 @@ use ruff_python_ast as ast;
 use ruff_python_ast::ParameterWithDefault;
 use ruff_python_semantic::analyze::function_type;
 use ruff_python_semantic::{Scope, ScopeKind, SemanticModel};
+use ruff_source_file::Locator;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -239,6 +240,7 @@ pub(crate) fn invalid_first_argument_name(
             parameters,
             checker.semantic(),
             function_type,
+            checker.locator(),
         )
     });
     diagnostics.push(diagnostic);
@@ -251,6 +253,7 @@ fn rename_parameter(
     parameters: &ast::Parameters,
     semantic: &SemanticModel<'_>,
     function_type: FunctionType,
+    locator: &Locator,
 ) -> Result<Option<Fix>> {
     // Don't fix if another parameter has the valid name.
     if parameters
@@ -272,6 +275,7 @@ fn rename_parameter(
         function_type.valid_first_argument_name(),
         scope,
         semantic,
+        locator,
     )?;
     Ok(Some(Fix::unsafe_edits(edit, rest)))
 }

crates/ruff_python_semantic/src/model.rs

@@ -1590,6 +1590,13 @@ impl<'a> SemanticModel<'a> {
             .intersects(SemanticModelFlags::COMPREHENSION_ASSIGNMENT)
     }
 
+    /// Return `true` if the model is visiting the r.h.s. of an `__all__` definition
+    /// (e.g. `"foo"` in `__all__ = ["foo"]`)
+    pub const fn in_dunder_all_definition(&self) -> bool {
+        self.flags
+            .intersects(SemanticModelFlags::DUNDER_ALL_DEFINITION)
+    }
+
     /// Return an iterator over all bindings shadowed by the given [`BindingId`], within the
     /// containing scope, and across scopes.
     pub fn shadowed_bindings(
@@ -1941,6 +1948,18 @@ bitflags! {
         /// ```
         const DOCSTRING = 1 << 21;
 
+        /// The model is visiting the r.h.s. of a module-level `__all__` definition.
+        ///
+        /// This could be any module-level statement that assigns or alters `__all__`,
+        /// for example:
+        /// ```python
+        /// __all__ = ["foo"]
+        /// __all__: str = ["foo"]
+        /// __all__ = ("bar",)
+        /// __all__ += ("baz,")
+        /// ```
+        const DUNDER_ALL_DEFINITION = 1 << 22;
+
         /// The context is in any type annotation.
         const ANNOTATION = Self::TYPING_ONLY_ANNOTATION.bits() | Self::RUNTIME_EVALUATED_ANNOTATION.bits() | Self::RUNTIME_REQUIRED_ANNOTATION.bits();
 

crates/ruff_python_semantic/src/reference.rs

@@ -93,6 +93,12 @@ impl ResolvedReference {
         self.flags
             .intersects(SemanticModelFlags::TYPE_CHECKING_BLOCK)
     }
+
+    /// Return `true` if the context is in the r.h.s. of an `__all__` definition.
+    pub const fn in_dunder_all_definition(&self) -> bool {
+        self.flags
+            .intersects(SemanticModelFlags::DUNDER_ALL_DEFINITION)
+    }
 }
 
 impl Ranged for ResolvedReference {