Correctly handle references in __all__ definitions when renaming sy…

…mbols in autofixes
astral-sh · Mar 22, 2024 · 7e7e801 · 7e7e801
1 parent d08be5d
commit 7e7e801
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 32 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI025_2.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI025_2.py
@@ -9,3 +9,6 @@
 
 if 1:
     __all__ += ["S" "e" "t"]
+
+if not False:
+    __all__ += ["Se" 't']
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI025_2.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI025_2.pyi
@@ -9,3 +9,6 @@ if True:
 
 if 1:
     __all__ += ["S" "e" "t"]
+
+if not False:
+    __all__ += ["Se" 't']
diff --git a/crates/ruff_linter/src/renamer.rs b/crates/ruff_linter/src/renamer.rs
@@ -4,10 +4,9 @@ use anyhow::{anyhow, Result};
 use itertools::Itertools;
 
 use ruff_diagnostics::Edit;
-use ruff_python_ast::str::raw_contents_range;
+use ruff_python_codegen::Stylist;
 use ruff_python_semantic::{Binding, BindingKind, Scope, ScopeId, SemanticModel};
-use ruff_source_file::Locator;
-use ruff_text_size::Ranged;
+use ruff_text_size::{Ranged, TextSize};
 
 pub(crate) struct Renamer;
 
@@ -106,7 +105,7 @@ impl Renamer {
         target: &str,
         scope: &Scope,
         semantic: &SemanticModel,
-        locator: &Locator,
+        stylist: &Stylist,
     ) -> Result<(Edit, Vec<Edit>)> {
         let mut edits = vec![];
 
@@ -134,7 +133,7 @@ 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, locator,
+            name, target, scope, semantic, stylist,
         ));
 
         // Find any scopes in which the symbol is referenced as `nonlocal` or `global`. For example,
@@ -166,7 +165,7 @@ impl Renamer {
         {
             let scope = &semantic.scopes[scope_id];
             edits.extend(Renamer::rename_in_scope(
-                name, target, scope, semantic, locator,
+                name, target, scope, semantic, stylist,
             ));
         }
 
@@ -187,7 +186,7 @@ impl Renamer {
         target: &str,
         scope: &Scope,
         semantic: &SemanticModel,
-        locator: &Locator,
+        stylist: &Stylist,
     ) -> Vec<Edit> {
         let mut edits = vec![];
 
@@ -210,21 +209,17 @@ impl Renamer {
                 // Rename the references to the binding.
                 edits.extend(binding.references().map(|reference_id| {
                     let reference = semantic.reference(reference_id);
-                    let range_to_replace = {
+                    let replacement = {
                         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()
+                            debug_assert!(!reference.range().is_empty());
+                            let quote = stylist.quote();
+                            format!("{quote}{target}{quote}")
                         } else {
-                            debug_assert!(locator.slice(reference.range()).contains(name));
-                            reference.range()
+                            debug_assert_eq!(TextSize::of(name), reference.range().len());
+                            target.to_string()
                         }
                     };
-                    Edit::range_replacement(target.to_string(), range_to_replace)
+                    Edit::range_replacement(replacement, reference.range())
                 }));
             }
         }

diff --git a/crates/ruff_linter/src/rules/flake8_import_conventions/rules/unconventional_import_alias.rs b/crates/ruff_linter/src/rules/flake8_import_conventions/rules/unconventional_import_alias.rs
@@ -64,9 +64,8 @@ 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(locator);
+    let name = binding.name(checker.locator());
     if binding.is_alias() && name == expected_alias {
         return None;
     }
@@ -82,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) =
-                    Renamer::rename(name, expected_alias, scope, checker.semantic(), locator)?;
+                let (edit, rest) = Renamer::rename(
+                    name,
+                    expected_alias,
+                    scope,
+                    checker.semantic(),
+                    checker.stylist(),
+                )?;
                 Ok(Fix::unsafe_edits(edit, rest))
             });
         }

diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/unaliased_collections_abc_set_import.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/unaliased_collections_abc_set_import.rs
@@ -68,8 +68,7 @@ pub(crate) fn unaliased_collections_abc_set_import(
         return None;
     }
 
-    let locator = checker.locator();
-    let name = binding.name(locator);
+    let name = binding.name(checker.locator());
     if name == "AbstractSet" {
         return None;
     }
@@ -78,8 +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(), locator)?;
+            let (edit, rest) = Renamer::rename(
+                name,
+                "AbstractSet",
+                scope,
+                checker.semantic(),
+                checker.stylist(),
+            )?;
             Ok(Fix::applicable_edits(
                 edit,
                 rest,

diff --git a/...rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI025_PYI025_2.py.snap b/...rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI025_PYI025_2.py.snap
@@ -23,8 +23,12 @@ PYI025_2.py:3:29: PYI025 [*] Use `from collections.abc import Set as AbstractSet
 6  6  | 
 7  7  | if True:
 8     |-    __all__ += [r'''Set''']
-   8  |+    __all__ += [r'''AbstractSet''']
+   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"]
diff --git a/...ules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI025_PYI025_2.pyi.snap b/...ules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI025_PYI025_2.pyi.snap
@@ -23,8 +23,12 @@ PYI025_2.pyi:3:29: PYI025 [*] Use `from collections.abc import Set as AbstractSe
 6  6  | 
 7  7  | if True:
 8     |-    __all__ += [r'''Set''']
-   8  |+    __all__ += [r'''AbstractSet''']
+   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"]
diff --git a/crates/ruff_linter/src/rules/pep8_naming/rules/invalid_first_argument_name.rs b/crates/ruff_linter/src/rules/pep8_naming/rules/invalid_first_argument_name.rs
@@ -4,9 +4,9 @@ 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_source_file::Locator;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -240,7 +240,7 @@ pub(crate) fn invalid_first_argument_name(
             parameters,
             checker.semantic(),
             function_type,
-            checker.locator(),
+            checker.stylist(),
         )
     });
     diagnostics.push(diagnostic);
@@ -253,7 +253,7 @@ fn rename_parameter(
     parameters: &ast::Parameters,
     semantic: &SemanticModel<'_>,
     function_type: FunctionType,
-    locator: &Locator,
+    stylist: &Stylist,
 ) -> Result<Option<Fix>> {
     // Don't fix if another parameter has the valid name.
     if parameters
@@ -275,7 +275,7 @@ fn rename_parameter(
         function_type.valid_first_argument_name(),
         scope,
         semantic,
-        locator,
+        stylist,
     )?;
     Ok(Some(Fix::unsafe_edits(edit, rest)))
 }