Add autofix for `Set`-to-`AbstractSet` rewrite using reference tracking (#5074)

## Summary

This PR enables autofix behavior for the `flake8-pyi` rule that asks you
to alias `Set` to `AbstractSet` when importing `collections.abc.Set`.
It's not the most important rule, but it's a good isolated test-case for
local symbol renaming.

The renaming algorithm is outlined in-detail in the `renamer.rs` module.
But to demonstrate the behavior, here's the diff when running this fix
over a complex file that exercises a few edge cases:

```diff
--- a/foo.pyi
+++ b/foo.pyi
@@ -1,16 +1,16 @@
 if True:
-    from collections.abc import Set
+    from collections.abc import Set as AbstractSet
 else:
-    Set = 1
+    AbstractSet = 1

-x: Set = set()
+x: AbstractSet = set()

-x: Set
+x: AbstractSet

-del Set
+del AbstractSet

 def f():
-    print(Set)
+    print(AbstractSet)

     def Set():
         pass
```

Making this work required resolving a bunch of edge cases in the
semantic model that were causing us to "lose track" of references. For
example, the above wasn't possible with our previous approach to
handling deletions (#5071). Similarly, the `x: Set` "delayed annotation"
tracking was enabled via #5070. And many of these edits would've failed
if we hadn't changed `BindingKind` to always match the identifier range
(#5090). So it's really the culmination of a bunch of changes over the
course of the week.

The main outstanding TODO is that this doesn't support `global` or
`nonlocal` usages. I'm going to take a look at that tonight, but I'm
comfortable merging this as-is.

Closes #1106.

Closes #5091.

crates/ruff/resources/test/fixtures/flake8_pyi/PYI025.py

@@ -1,19 +1,19 @@
-from collections.abc import Set as AbstractSet  # Ok
+def f():
+    from collections.abc import Set as AbstractSet  # Ok
 
 
-from collections.abc import Set  # Ok
+def f():
+    from collections.abc import Container, Sized, Set as AbstractSet, ValuesView  # Ok
 
 
-from collections.abc import (
+def f():
-    Container,
+    from collections.abc import Set  # PYI025
-    Sized,
-    Set,  # Ok
-    ValuesView
-)
 
-from collections.abc import (
+
-    Container,
+def f():
-    Sized,
+    from collections.abc import Container, Sized, Set, ValuesView  # PYI025
-    Set as AbstractSet,  # Ok
+
-    ValuesView
+    GLOBAL: Set[int] = set()
-)
+
+    class Class:
+        member: Set[int]

crates/ruff/resources/test/fixtures/flake8_pyi/PYI025.pyi

@@ -1,19 +1,30 @@
-from collections.abc import Set as AbstractSet  # Ok
+def f():
+    from collections.abc import Set as AbstractSet  # Ok
 
+def f():
+    from collections.abc import Container, Sized, Set as AbstractSet, ValuesView  # Ok
 
-from collections.abc import Set  # PYI025
+def f():
+    from collections.abc import Set  # PYI025
 
+def f():
+    from collections.abc import Container, Sized, Set, ValuesView  # PYI025
 
-from collections.abc import (
+def f():
-    Container,
+    if True:
-    Sized,
+        from collections.abc import Set
-    Set,  # PYI025
+    else:
-    ValuesView
+        Set = 1
-)
 
-from collections.abc import (
+    x: Set = set()
-    Container,
+
-    Sized,
+    x: Set
-    Set as AbstractSet,
+
-    ValuesView # Ok
+    del Set
-)
+
+    def f():
+        print(Set)
+
+        def Set():
+            pass
+        print(Set)

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

@@ -817,24 +817,28 @@ where
                             BindingKind::SubmoduleImportation(SubmoduleImportation {
                                 qualified_name,
                             }),
-                            BindingFlags::empty(),
+                            BindingFlags::EXTERNAL,
                         );
                     } else {
+                        let mut flags = BindingFlags::EXTERNAL;
+                        if alias.asname.is_some() {
+                            flags |= BindingFlags::ALIAS;
+                        }
+                        if alias
+                            .asname
+                            .as_ref()
+                            .map_or(false, |asname| asname == &alias.name)
+                        {
+                            flags |= BindingFlags::EXPLICIT_EXPORT;
+                        }
+
                         let name = alias.asname.as_ref().unwrap_or(&alias.name);
                         let qualified_name = &alias.name;
                         self.add_binding(
                             name,
                             alias.identifier(self.locator),
                             BindingKind::Importation(Importation { qualified_name }),
-                            if alias
+                            flags,
-                                .asname
-                                .as_ref()
-                                .map_or(false, |asname| asname == &alias.name)
-                            {
-                                BindingFlags::EXPLICIT_EXPORT
-                            } else {
-                                BindingFlags::empty()
-                            },
                         );
 
                         if let Some(asname) = &alias.asname {
@@ -1052,9 +1056,6 @@ where
                 }
 
                 if self.is_stub {
-                    if self.enabled(Rule::UnaliasedCollectionsAbcSetImport) {
-                        flake8_pyi::rules::unaliased_collections_abc_set_import(self, import_from);
-                    }
                     if self.enabled(Rule::FutureAnnotationsInStub) {
                         flake8_pyi::rules::from_future_import(self, import_from);
                     }
@@ -1116,6 +1117,18 @@ where
                             }
                         }
 
+                        let mut flags = BindingFlags::EXTERNAL;
+                        if alias.asname.is_some() {
+                            flags |= BindingFlags::ALIAS;
+                        }
+                        if alias
+                            .asname
+                            .as_ref()
+                            .map_or(false, |asname| asname == &alias.name)
+                        {
+                            flags |= BindingFlags::EXPLICIT_EXPORT;
+                        }
+
                         // Given `from foo import bar`, `name` would be "bar" and `qualified_name` would
                         // be "foo.bar". Given `from foo import bar as baz`, `name` would be "baz"
                         // and `qualified_name` would be "foo.bar".
@@ -1126,15 +1139,7 @@ where
                             name,
                             alias.identifier(self.locator),
                             BindingKind::FromImportation(FromImportation { qualified_name }),
-                            if alias
+                            flags,
-                                .asname
-                                .as_ref()
-                                .map_or(false, |asname| asname == &alias.name)
-                            {
-                                BindingFlags::EXPLICIT_EXPORT
-                            } else {
-                                BindingFlags::empty()
-                            },
                         );
                     }
                     if self.enabled(Rule::RelativeImports) {
@@ -4711,23 +4716,18 @@ impl<'a> Checker<'a> {
     }
 
     fn check_dead_scopes(&mut self) {
-        let enforce_typing_imports = !self.is_stub
+        if !self.any_enabled(&[
-            && self.any_enabled(&[
+            Rule::UnusedImport,
             Rule::GlobalVariableNotAssigned,
+            Rule::UndefinedLocalWithImportStarUsage,
+            Rule::RedefinedWhileUnused,
             Rule::RuntimeImportInTypeCheckingBlock,
             Rule::TypingOnlyFirstPartyImport,
             Rule::TypingOnlyThirdPartyImport,
             Rule::TypingOnlyStandardLibraryImport,
-            ]);
-
-        if !(enforce_typing_imports
-            || self.any_enabled(&[
-                Rule::UnusedImport,
-                Rule::UndefinedLocalWithImportStarUsage,
-                Rule::RedefinedWhileUnused,
             Rule::UndefinedExport,
-            ]))
+            Rule::UnaliasedCollectionsAbcSetImport,
-        {
+        ]) {
             return;
         }
 
@@ -4757,6 +4757,16 @@ impl<'a> Checker<'a> {
         // Identify any valid runtime imports. If a module is imported at runtime, and
         // used at runtime, then by default, we avoid flagging any other
         // imports from that model as typing-only.
+        let enforce_typing_imports = if self.is_stub {
+            false
+        } else {
+            self.any_enabled(&[
+                Rule::RuntimeImportInTypeCheckingBlock,
+                Rule::TypingOnlyFirstPartyImport,
+                Rule::TypingOnlyThirdPartyImport,
+                Rule::TypingOnlyStandardLibraryImport,
+            ])
+        };
         let runtime_imports: Vec<Vec<&Binding>> = if enforce_typing_imports {
             if self.settings.flake8_type_checking.strict {
                 vec![]
@@ -4907,6 +4917,16 @@ impl<'a> Checker<'a> {
             if self.enabled(Rule::UnusedImport) {
                 pyflakes::rules::unused_import(self, scope, &mut diagnostics);
             }
+
+            if self.is_stub {
+                if self.enabled(Rule::UnaliasedCollectionsAbcSetImport) {
+                    flake8_pyi::rules::unaliased_collections_abc_set_import(
+                        self,
+                        scope,
+                        &mut diagnostics,
+                    );
+                }
+            }
         }
         self.diagnostics.extend(diagnostics);
     }

crates/ruff/src/importer/mod.rs

@@ -1,4 +1,7 @@
-//! Add and modify import statements to make module members available during fix execution.
+//! Code modification struct to add and modify import statements.
+//!
+//! Enables rules to make module members available (that may be not yet be imported) during fix
+//! execution.
 
 use std::error::Error;
 

crates/ruff/src/lib.rs

@@ -29,6 +29,7 @@ mod noqa;
 pub mod packaging;
 pub mod pyproject_toml;
 pub mod registry;
+mod renamer;
 pub mod resolver;
 mod rule_redirects;
 mod rule_selector;

crates/ruff/src/renamer.rs

@@ -0,0 +1,176 @@
+//! Code modification struct to support symbol renaming within a scope.
+
+use anyhow::{anyhow, Result};
+
+use ruff_diagnostics::Edit;
+use ruff_python_semantic::{Binding, BindingKind, Scope, SemanticModel};
+
+pub(crate) struct Renamer;
+
+impl Renamer {
+    /// Rename a symbol (from `name` to `target`) within a [`Scope`].
+    ///
+    /// ## How it works
+    ///
+    /// The renaming algorithm is as follows:
+    ///
+    /// 1. Start with the first [`Binding`] in the scope, for the given name. For example, in the
+    ///    following snippet, we'd start by examining the `x = 1` binding:
+    ///
+    /// ```python
+    /// if True:
+    ///     x = 1
+    ///     print(x)
+    /// else:
+    ///     x = 2
+    ///     print(x)
+    ///
+    /// print(x)
+    /// ```
+    ///
+    /// 2. Rename the [`Binding`]. In most cases, this is a simple replacement. For example,
+    ///    renaming `x` to `y` above would require replacing `x = 1` with `y = 1`. After the
+    ///    first replacement in the snippet above, we'd have:
+    ///
+    /// ```python
+    /// if True:
+    ///     y = 1
+    ///     print(x)
+    /// else:
+    ///     x = 2
+    ///     print(x)
+    ///
+    /// print(x)
+    /// ```
+    ///
+    ///    Note that, when renaming imports, we need to instead rename (or add) an alias. For
+    ///    example, to rename `pandas` to `pd`, we may need to rewrite `import pandas` to
+    ///    `import pandas as pd`, rather than `import pd`.
+    ///
+    /// 3. Rename every reference to the [`Binding`]. For example, renaming the references to the
+    ///    `x = 1` binding above would give us:
+    ///
+    /// ```python
+    /// if True:
+    ///     y = 1
+    ///     print(y)
+    /// else:
+    ///     x = 2
+    ///     print(x)
+    ///
+    /// print(x)
+    /// ```
+    ///
+    /// 4. Rename every delayed annotation. (See [`SemanticModel::delayed_annotations`].)
+    ///
+    /// 5. Repeat the above process for every [`Binding`] in the scope with the given name.
+    ///    After renaming the `x = 2` binding, we'd have:
+    ///
+    /// ```python
+    /// if True:
+    ///     y = 1
+    ///     print(y)
+    /// else:
+    ///     y = 2
+    ///     print(y)
+    ///
+    /// print(y)
+    /// ```
+    ///
+    /// ## Limitations
+    ///
+    /// `global` and `nonlocal` declarations are not yet supported.
+    ///
+    /// `global` and `nonlocal` declarations add some additional complexity. If we're renaming a
+    /// name that's declared as `global` or `nonlocal` in a child scope, we need to rename the name
+    /// in that scope too, repeating the above process.
+    ///
+    /// If we're renaming a name that's declared as `global` or `nonlocal` in the current scope,
+    /// then we need to identify the scope in which the name is declared, and perform the rename
+    /// in that scope instead (which will in turn trigger the above process on the current scope).
+    pub(crate) fn rename(
+        name: &str,
+        target: &str,
+        scope: &Scope,
+        semantic: &SemanticModel,
+    ) -> Result<(Edit, Vec<Edit>)> {
+        let mut edits = vec![];
+
+        // Iterate over every binding to the name in the scope.
+        for binding_id in scope.get_all(name) {
+            let binding = semantic.binding(binding_id);
+
+            // Rename the binding.
+            if let Some(edit) = Renamer::rename_binding(binding, name, target) {
+                edits.push(edit);
+
+                // Rename any delayed annotations.
+                if let Some(annotations) = semantic.delayed_annotations(binding_id) {
+                    edits.extend(annotations.iter().filter_map(|annotation_id| {
+                        let annotation = semantic.binding(*annotation_id);
+                        Renamer::rename_binding(annotation, name, target)
+                    }));
+                }
+
+                // 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())
+                }));
+            }
+        }
+
+        // Deduplicate any edits. In some cases, a reference can be both a read _and_ a write. For
+        // example, `x += 1` is both a read of and a write to `x`.
+        edits.sort();
+        edits.dedup();
+
+        let edit = edits
+            .pop()
+            .ok_or(anyhow!("Unable to rename any references to `{name}`"))?;
+        Ok((edit, edits))
+    }
+
+    /// Rename a [`Binding`] reference.
+    fn rename_binding(binding: &Binding, name: &str, target: &str) -> Option<Edit> {
+        match &binding.kind {
+            BindingKind::Importation(_) | BindingKind::FromImportation(_) => {
+                if binding.is_alias() {
+                    // Ex) Rename `import pandas as alias` to `import pandas as pd`.
+                    Some(Edit::range_replacement(target.to_string(), binding.range))
+                } else {
+                    // Ex) Rename `import pandas` to `import pandas as pd`.
+                    Some(Edit::range_replacement(
+                        format!("{name} as {target}"),
+                        binding.range,
+                    ))
+                }
+            }
+            BindingKind::SubmoduleImportation(import) => {
+                // Ex) Rename `import pandas.core` to `import pandas as pd`.
+                let module_name = import.qualified_name.split('.').next().unwrap();
+                Some(Edit::range_replacement(
+                    format!("{module_name} as {target}"),
+                    binding.range,
+                ))
+            }
+            // Avoid renaming builtins and other "special" bindings.
+            BindingKind::FutureImportation | BindingKind::Builtin | BindingKind::Export(_) => None,
+            // By default, replace the binding's name with the target name.
+            BindingKind::Annotation
+            | BindingKind::Argument
+            | BindingKind::NamedExprAssignment
+            | BindingKind::UnpackedAssignment
+            | BindingKind::Assignment
+            | BindingKind::LoopVar
+            | BindingKind::Global
+            | BindingKind::Nonlocal
+            | BindingKind::ClassDefinition
+            | BindingKind::FunctionDefinition
+            | BindingKind::Deletion
+            | BindingKind::UnboundException => {
+                Some(Edit::range_replacement(target.to_string(), binding.range))
+            }
+        }
+    }
+}

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

@@ -1,9 +1,10 @@
-use rustpython_parser::ast::StmtImportFrom;
+use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation};
-
-use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::{BindingKind, FromImportation, Scope};
 
 use crate::checkers::ast::Checker;
+use crate::registry::AsRule;
+use crate::renamer::Renamer;
 
 /// ## What it does
 /// Checks for `from collections.abc import Set` imports that do not alias
@@ -30,6 +31,8 @@ use crate::checkers::ast::Checker;
 pub struct UnaliasedCollectionsAbcSetImport;
 
 impl Violation for UnaliasedCollectionsAbcSetImport {
+    const AUTOFIX: AutofixKind = AutofixKind::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         format!(
@@ -43,20 +46,33 @@ impl Violation for UnaliasedCollectionsAbcSetImport {
 }
 
 /// PYI025
-pub(crate) fn unaliased_collections_abc_set_import(checker: &mut Checker, stmt: &StmtImportFrom) {
+pub(crate) fn unaliased_collections_abc_set_import(
-    let Some(module_id) = &stmt.module else {
+    checker: &Checker,
-        return;
+    scope: &Scope,
+    diagnostics: &mut Vec<Diagnostic>,
+) {
+    for (name, binding_id) in scope.all_bindings() {
+        let binding = checker.semantic().binding(binding_id);
+        let BindingKind::FromImportation(FromImportation { qualified_name }) = &binding.kind else {
+            continue;
         };
-    if module_id.as_str() != "collections.abc" {
+        if qualified_name.as_str() != "collections.abc.Set" {
-        return;
+            continue;
+        }
+        if name == "AbstractSet" {
+            continue;
         }
 
-    for name in &stmt.names {
+        let mut diagnostic = Diagnostic::new(UnaliasedCollectionsAbcSetImport, binding.range);
-        if name.name.as_str() == "Set" && name.asname.is_none() {
+        if checker.patch(diagnostic.kind.rule()) {
-            checker.diagnostics.push(Diagnostic::new(
+            if checker.semantic().is_available("AbstractSet") {
-                UnaliasedCollectionsAbcSetImport,
+                diagnostic.try_set_fix(|| {
-                name.range,
+                    let (edit, rest) =
-            ));
+                        Renamer::rename(name, "AbstractSet", scope, checker.semantic())?;
+                    Ok(Fix::suggested_edits(edit, rest))
+                });
             }
         }
+        diagnostics.push(diagnostic);
+    }
 }

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI025_PYI025.pyi.snap

@@ -1,22 +1,81 @@
 ---
 source: crates/ruff/src/rules/flake8_pyi/mod.rs
 ---
-PYI025.pyi:4:29: PYI025 Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+PYI025.pyi:8:33: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
    |
-4 | from collections.abc import Set  # PYI025
+ 7 | def f():
+ 8 |     from collections.abc import Set  # PYI025
    |                                 ^^^ PYI025
+ 9 | 
+10 | def f():
    |
    = help: Alias `Set` to `AbstractSet`
 
-PYI025.pyi:10:5: PYI025 Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+ℹ Suggested fix
+5 5 |     from collections.abc import Container, Sized, Set as AbstractSet, ValuesView  # Ok
+6 6 | 
+7 7 | def f():
+8   |-    from collections.abc import Set  # PYI025
+  8 |+    from collections.abc import Set as AbstractSet  # PYI025
+9 9 | 
+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
    |
- 8 |     Container,
+10 | def f():
- 9 |     Sized,
+11 |     from collections.abc import Container, Sized, Set, ValuesView  # PYI025
-10 |     Set,  # PYI025
    |                                                   ^^^ PYI025
-11 |     ValuesView
+12 | 
-12 | )
+13 | def f():
    |
    = help: Alias `Set` to `AbstractSet`
 
+ℹ Suggested fix
+8  8  |     from collections.abc import Set  # PYI025
+9  9  | 
+10 10 | def f():
+11    |-    from collections.abc import Container, Sized, Set, ValuesView  # PYI025
+   11 |+    from collections.abc import Container, Sized, Set as AbstractSet, ValuesView  # PYI025
+12 12 | 
+13 13 | def f():
+14 14 |     if True:
+
+PYI025.pyi:15:37: PYI025 [*] Use `from collections.abc import Set as AbstractSet` to avoid confusion with the `set` builtin
+   |
+13 | def f():
+14 |     if True:
+15 |         from collections.abc import Set
+   |                                     ^^^ PYI025
+16 |     else:
+17 |         Set = 1
+   |
+   = help: Alias `Set` to `AbstractSet`
+
+ℹ Suggested fix
+12 12 | 
+13 13 | def f():
+14 14 |     if True:
+15    |-        from collections.abc import Set
+   15 |+        from collections.abc import Set as AbstractSet
+16 16 |     else:
+17    |-        Set = 1
+   17 |+        AbstractSet = 1
+18 18 | 
+19    |-    x: Set = set()
+   19 |+    x: AbstractSet = set()
+20 20 | 
+21    |-    x: Set
+   21 |+    x: AbstractSet
+22 22 | 
+23    |-    del Set
+   23 |+    del AbstractSet
+24 24 | 
+25 25 |     def f():
+26    |-        print(Set)
+   26 |+        print(AbstractSet)
+27 27 | 
+28 28 |         def Set():
+29 29 |             pass
+
 

crates/ruff_python_semantic/src/binding.rs

@@ -44,6 +44,18 @@ impl<'a> Binding<'a> {
         self.flags.contains(BindingFlags::EXPLICIT_EXPORT)
     }
 
+    /// Return `true` if this [`Binding`] represents an external symbol
+    /// (e.g., `FastAPI` in `from fastapi import FastAPI`).
+    pub const fn is_external(&self) -> bool {
+        self.flags.contains(BindingFlags::EXTERNAL)
+    }
+
+    /// Return `true` if this [`Binding`] represents an aliased symbol
+    /// (e.g., `app` in `from fastapi import FastAPI as app`).
+    pub const fn is_alias(&self) -> bool {
+        self.flags.contains(BindingFlags::ALIAS)
+    }
+
     /// Return `true` if this [`Binding`] represents an unbound variable
     /// (e.g., `x` in `x = 1; del x`).
     pub const fn is_unbound(&self) -> bool {
@@ -161,9 +173,25 @@ bitflags! {
         ///
         /// For example, the binding could be `FastAPI` in:
         /// ```python
-        /// import FastAPI as FastAPI
+        /// from fastapi import FastAPI as FastAPI
         /// ```
         const EXPLICIT_EXPORT = 1 << 0;
+
+        /// The binding represents an external symbol, like an import or a builtin.
+        ///
+        /// For example, the binding could be `FastAPI` in:
+        /// ```python
+        /// from fastapi import FastAPI
+        /// ```
+        const EXTERNAL = 1 << 1;
+
+        /// The binding is an aliased symbol.
+        ///
+        /// For example, the binding could be `app` in:
+        /// ```python
+        /// from fastapi import FastAPI as app
+        /// ```
+        const ALIAS = 1 << 2;
     }
 }