Correctly handle references in __all__ definitions when renaming symbols in autofixes (#10527)

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

@@ -0,0 +1,14 @@
+"""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"]
+
+if not False:
+    __all__ += ["Se" 't']

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

@@ -0,0 +1,14 @@
+"""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"]
+
+if not False:
+    __all__ += ["Se" 't']

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

@@ -2114,9 +2114,11 @@ impl<'a> Checker<'a> {
         for export in exports {
             let (name, range) = (export.name(), export.range());
             if let Some(binding_id) = self.semantic.global_scope().get(name) {
+                self.semantic.flags |= SemanticModelFlags::DUNDER_ALL_DEFINITION;
                 // Mark anything referenced in `__all__` as used.
                 self.semantic
                     .add_global_reference(binding_id, ExprContext::Load, range);
+                self.semantic.flags -= SemanticModelFlags::DUNDER_ALL_DEFINITION;
             } else {
                 if self.semantic.global_scope().uses_star_imports() {
                     if self.enabled(Rule::UndefinedLocalWithImportStarUsage) {

crates/ruff_linter/src/renamer.rs

@@ -4,8 +4,9 @@ use anyhow::{anyhow, Result};
 use itertools::Itertools;
 
 use ruff_diagnostics::Edit;
+use ruff_python_codegen::Stylist;
 use ruff_python_semantic::{Binding, BindingKind, Scope, ScopeId, SemanticModel};
-use ruff_text_size::Ranged;
+use ruff_text_size::{Ranged, TextSize};
 
 pub(crate) struct Renamer;
 
@@ -104,6 +105,7 @@ impl Renamer {
         target: &str,
         scope: &Scope,
         semantic: &SemanticModel,
+        stylist: &Stylist,
     ) -> Result<(Edit, Vec<Edit>)> {
         let mut edits = vec![];
 
@@ -130,7 +132,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, stylist,
+        ));
 
         // Find any scopes in which the symbol is referenced as `nonlocal` or `global`. For example,
         // given:
@@ -160,7 +164,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, stylist,
+            ));
         }
 
         // Deduplicate any edits.
@@ -180,6 +186,7 @@ impl Renamer {
         target: &str,
         scope: &Scope,
         semantic: &SemanticModel,
+        stylist: &Stylist,
     ) -> Vec<Edit> {
         let mut edits = vec![];
 
@@ -202,7 +209,17 @@ 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 replacement = {
+                        if reference.in_dunder_all_definition() {
+                            debug_assert!(!reference.range().is_empty());
+                            let quote = stylist.quote();
+                            format!("{quote}{target}{quote}")
+                        } else {
+                            debug_assert_eq!(TextSize::of(name), reference.range().len());
+                            target.to_string()
+                        }
+                    };
+                    Edit::range_replacement(replacement, reference.range())
                 }));
             }
         }

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

@@ -81,8 +81,13 @@ pub(crate) fn unconventional_import_alias(
         if checker.semantic().is_available(expected_alias) {
             diagnostic.try_set_fix(|| {
                 let scope = &checker.semantic().scopes[binding.scope];
-                let (edit, rest) =
+                let (edit, rest) = Renamer::rename(
-                    Renamer::rename(name, expected_alias, scope, checker.semantic())?;
+                    name,
+                    expected_alias,
+                    scope,
+                    checker.semantic(),
+                    checker.stylist(),
+                )?;
                 Ok(Fix::unsafe_edits(edit, rest))
             });
         }

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

@@ -79,8 +79,10 @@ mod tests {
     #[test_case(Rule::TypeCommentInStub, Path::new("PYI033.pyi"))]
     #[test_case(Rule::TypedArgumentDefaultInStub, Path::new("PYI011.py"))]
     #[test_case(Rule::TypedArgumentDefaultInStub, Path::new("PYI011.pyi"))]
-    #[test_case(Rule::UnaliasedCollectionsAbcSetImport, Path::new("PYI025.py"))]
+    #[test_case(Rule::UnaliasedCollectionsAbcSetImport, Path::new("PYI025_1.py"))]
-    #[test_case(Rule::UnaliasedCollectionsAbcSetImport, Path::new("PYI025.pyi"))]
+    #[test_case(Rule::UnaliasedCollectionsAbcSetImport, Path::new("PYI025_1.pyi"))]
+    #[test_case(Rule::UnaliasedCollectionsAbcSetImport, Path::new("PYI025_2.py"))]
+    #[test_case(Rule::UnaliasedCollectionsAbcSetImport, Path::new("PYI025_2.pyi"))]
     #[test_case(Rule::UnannotatedAssignmentInStub, Path::new("PYI052.py"))]
     #[test_case(Rule::UnannotatedAssignmentInStub, Path::new("PYI052.pyi"))]
     #[test_case(Rule::UnassignedSpecialVariableInStub, Path::new("PYI035.py"))]

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

@@ -77,7 +77,13 @@ 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(),
+                checker.stylist(),
+            )?;
             Ok(Fix::applicable_edits(
                 edit,
                 rest,

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

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
 ---
-PYI025.py:10:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025_1.py:10:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
  9 | def f():
 10 |     from collections.abc import Set  # PYI025
@@ -19,7 +19,7 @@ PYI025.py:10:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet`
 12 12 | 
 13 13 | def f():
 
-PYI025.py:14:51: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025_1.py:14:51: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
 13 | def f():
 14 |     from collections.abc import Container, Sized, Set, ValuesView  # PYI025
@@ -42,5 +42,3 @@ PYI025.py:14:51: PYI025 [*] Use `from collections.abc import Set as AbstractSet`
 18 18 |     class Class:
 19    |-        member: Set[int]
    19 |+        member: AbstractSet[int]
-
-

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

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
 ---
-PYI025.pyi:8:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025_1.pyi:8:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
  7 | def f():
  8 |     from collections.abc import Set  # PYI025
@@ -21,7 +21,7 @@ PYI025.pyi:8:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet`
 10 10 | def f():
 11 11 |     from collections.abc import Container, Sized, Set, ValuesView  # PYI025
 
-PYI025.pyi:11:51: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025_1.pyi:11:51: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
 10 | def f():
 11 |     from collections.abc import Container, Sized, Set, ValuesView  # PYI025
@@ -41,7 +41,7 @@ PYI025.pyi:11:51: PYI025 [*] Use `from collections.abc import Set as AbstractSet
 13 13 | def f():
 14 14 |     """Test: local symbol renaming."""
 
-PYI025.pyi:16:37: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025_1.pyi:16:37: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
 14 |     """Test: local symbol renaming."""
 15 |     if True:
@@ -76,7 +76,7 @@ PYI025.pyi:16:37: PYI025 [*] Use `from collections.abc import Set as AbstractSet
 29 29 |         def Set():
 30 30 |             pass
 
-PYI025.pyi:33:29: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025_1.pyi:33:29: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
 31 |         print(Set)
 32 | 
@@ -119,7 +119,7 @@ PYI025.pyi:33:29: PYI025 [*] Use `from collections.abc import Set as AbstractSet
 42 42 | def f():
 43 43 |     """Test: nonlocal symbol renaming."""
 
-PYI025.pyi:44:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025_1.pyi:44:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
 42 | def f():
 43 |     """Test: nonlocal symbol renaming."""
@@ -145,5 +145,3 @@ PYI025.pyi:44:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet
 50    |-        print(Set)
    49 |+        AbstractSet = 1
    50 |+        print(AbstractSet)
-
-

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

@@ -0,0 +1,34 @@
+---
+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__ += ["AbstractSet"]
+9  9  | 
+10 10 | if 1:
+11    |-    __all__ += ["S" "e" "t"]
+   11 |+    __all__ += ["AbstractSet"]
+12 12 | 
+13 13 | if not False:
+14    |-    __all__ += ["Se" 't']
+   14 |+    __all__ += ["AbstractSet"]

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

@@ -0,0 +1,34 @@
+---
+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__ += ["AbstractSet"]
+9  9  | 
+10 10 | if 1:
+11    |-    __all__ += ["S" "e" "t"]
+   11 |+    __all__ += ["AbstractSet"]
+12 12 | 
+13 13 | if not False:
+14    |-    __all__ += ["Se" 't']
+   14 |+    __all__ += ["AbstractSet"]

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

@@ -4,6 +4,7 @@ use ruff_diagnostics::{Diagnostic, DiagnosticKind, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast as ast;
 use ruff_python_ast::ParameterWithDefault;
+use ruff_python_codegen::Stylist;
 use ruff_python_semantic::analyze::function_type;
 use ruff_python_semantic::{Scope, ScopeKind, SemanticModel};
 use ruff_text_size::Ranged;
@@ -239,6 +240,7 @@ pub(crate) fn invalid_first_argument_name(
             parameters,
             checker.semantic(),
             function_type,
+            checker.stylist(),
         )
     });
     diagnostics.push(diagnostic);
@@ -251,6 +253,7 @@ fn rename_parameter(
     parameters: &ast::Parameters,
     semantic: &SemanticModel<'_>,
     function_type: FunctionType,
+    stylist: &Stylist,
 ) -> 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,
+        stylist,
     )?;
     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 {