Allow removal of typing from exempt-modules (#9214)

## Summary

If you remove `typing` from `exempt-modules`, we tend to panic, since we
try to add `TYPE_CHECKING` to `from typing import ...` statements while
concurrently attempting to remove other members from that import. This
PR adds special-casing for typing imports to avoid such panics.

Closes https://github.com/astral-sh/ruff/issues/5331
Closes https://github.com/astral-sh/ruff/issues/9196.
Closes https://github.com/astral-sh/ruff/issues/9197.

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/exempt_type_checking_1.py

@@ -0,0 +1,7 @@
+"""Add `TYPE_CHECKING` to an existing `typing` import. Another member is moved."""
+
+from __future__ import annotations
+
+from typing import Final
+
+Const: Final[dict] = {}

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/exempt_type_checking_2.py

@@ -0,0 +1,7 @@
+"""Using `TYPE_CHECKING` from an existing `typing` import. Another member is moved."""
+
+from __future__ import annotations
+
+from typing import Final, TYPE_CHECKING
+
+Const: Final[dict] = {}

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/exempt_type_checking_3.py

@@ -0,0 +1,7 @@
+"""Using `TYPE_CHECKING` from an existing `typing` import. Another member is moved."""
+
+from __future__ import annotations
+
+from typing import Final, Mapping
+
+Const: Final[dict] = {}

crates/ruff_linter/src/importer/mod.rs

@@ -13,7 +13,7 @@ use ruff_text_size::{Ranged, TextSize};
 use ruff_diagnostics::Edit;
 use ruff_python_ast::imports::{AnyImport, Import, ImportFrom};
 use ruff_python_codegen::Stylist;
-use ruff_python_semantic::SemanticModel;
+use ruff_python_semantic::{ImportedName, SemanticModel};
 use ruff_python_trivia::textwrap::indent;
 use ruff_source_file::Locator;
 
@@ -132,7 +132,48 @@ impl<'a> Importer<'a> {
         )?;
 
         // Import the `TYPE_CHECKING` symbol from the typing module.
-        let (type_checking_edit, type_checking) = self.get_or_import_type_checking(at, semantic)?;
+        let (type_checking_edit, type_checking) =
+            if let Some(type_checking) = Self::find_type_checking(at, semantic)? {
+                // Special-case: if the `TYPE_CHECKING` symbol is imported as part of the same
+                // statement that we're modifying, avoid adding a no-op edit. For example, here,
+                // the `TYPE_CHECKING` no-op edit would overlap with the edit to remove `Final`
+                // from the import:
+                // ```python
+                // from __future__ import annotations
+                //
+                // from typing import Final, TYPE_CHECKING
+                //
+                // Const: Final[dict] = {}
+                // ```
+                let edit = if type_checking.statement(semantic) == import.statement {
+                    None
+                } else {
+                    Some(Edit::range_replacement(
+                        self.locator.slice(type_checking.range()).to_string(),
+                        type_checking.range(),
+                    ))
+                };
+                (edit, type_checking.into_name())
+            } else {
+                // Special-case: if the `TYPE_CHECKING` symbol would be added to the same import
+                // we're modifying, import it as a separate import statement. For example, here,
+                // we're concurrently removing `Final` and adding `TYPE_CHECKING`, so it's easier to
+                // use a separate import statement:
+                // ```python
+                // from __future__ import annotations
+                //
+                // from typing import Final
+                //
+                // Const: Final[dict] = {}
+                // ```
+                let (edit, name) = self.import_symbol(
+                    &ImportRequest::import_from("typing", "TYPE_CHECKING"),
+                    at,
+                    Some(import.statement),
+                    semantic,
+                )?;
+                (Some(edit), name)
+            };
 
         // Add the import to a `TYPE_CHECKING` block.
         let add_import_edit = if let Some(block) = self.preceding_type_checking_block(at) {
@@ -157,28 +198,21 @@ impl<'a> Importer<'a> {
         })
     }
 
-    /// Generate an [`Edit`] to reference `typing.TYPE_CHECKING`. Returns the [`Edit`] necessary to
+    /// Find a reference to `typing.TYPE_CHECKING`.
-    /// make the symbol available in the current scope along with the bound name of the symbol.
+    fn find_type_checking(
-    fn get_or_import_type_checking(
-        &self,
         at: TextSize,
         semantic: &SemanticModel,
-    ) -> Result<(Edit, String), ResolutionError> {
+    ) -> Result<Option<ImportedName>, ResolutionError> {
         for module in semantic.typing_modules() {
-            if let Some((edit, name)) = self.get_symbol(
+            if let Some(imported_name) = Self::find_symbol(
                 &ImportRequest::import_from(module, "TYPE_CHECKING"),
                 at,
                 semantic,
             )? {
-                return Ok((edit, name));
+                return Ok(Some(imported_name));
             }
         }
-
+        Ok(None)
-        self.import_symbol(
-            &ImportRequest::import_from("typing", "TYPE_CHECKING"),
-            at,
-            semantic,
-        )
     }
 
     /// Generate an [`Edit`] to reference the given symbol. Returns the [`Edit`] necessary to make
@@ -192,16 +226,15 @@ impl<'a> Importer<'a> {
         semantic: &SemanticModel,
     ) -> Result<(Edit, String), ResolutionError> {
         self.get_symbol(symbol, at, semantic)?
-            .map_or_else(|| self.import_symbol(symbol, at, semantic), Ok)
+            .map_or_else(|| self.import_symbol(symbol, at, None, semantic), Ok)
     }
 
-    /// Return an [`Edit`] to reference an existing symbol, if it's present in the given [`SemanticModel`].
+    /// Return the [`ImportedName`] to for existing symbol, if it's present in the given [`SemanticModel`].
-    fn get_symbol(
+    fn find_symbol(
-        &self,
         symbol: &ImportRequest,
         at: TextSize,
         semantic: &SemanticModel,
-    ) -> Result<Option<(Edit, String)>, ResolutionError> {
+    ) -> Result<Option<ImportedName>, ResolutionError> {
         // If the symbol is already available in the current scope, use it.
         let Some(imported_name) =
             semantic.resolve_qualified_import_name(symbol.module, symbol.member)
@@ -226,6 +259,21 @@ impl<'a> Importer<'a> {
             return Err(ResolutionError::IncompatibleContext);
         }
 
+        Ok(Some(imported_name))
+    }
+
+    /// Return an [`Edit`] to reference an existing symbol, if it's present in the given [`SemanticModel`].
+    fn get_symbol(
+        &self,
+        symbol: &ImportRequest,
+        at: TextSize,
+        semantic: &SemanticModel,
+    ) -> Result<Option<(Edit, String)>, ResolutionError> {
+        // Find the symbol in the current scope.
+        let Some(imported_name) = Self::find_symbol(symbol, at, semantic)? else {
+            return Ok(None);
+        };
+
         // We also add a no-op edit to force conflicts with any other fixes that might try to
         // remove the import. Consider:
         //
@@ -259,9 +307,13 @@ impl<'a> Importer<'a> {
         &self,
         symbol: &ImportRequest,
         at: TextSize,
+        except: Option<&Stmt>,
         semantic: &SemanticModel,
     ) -> Result<(Edit, String), ResolutionError> {
-        if let Some(stmt) = self.find_import_from(symbol.module, at) {
+        if let Some(stmt) = self
+            .find_import_from(symbol.module, at)
+            .filter(|stmt| except != Some(stmt))
+        {
             // Case 1: `from functools import lru_cache` is in scope, and we're trying to reference
             // `functools.cache`; thus, we add `cache` to the import, and return `"cache"` as the
             // bound name.
@@ -423,14 +475,18 @@ impl RuntimeImportEdit {
 #[derive(Debug)]
 pub(crate) struct TypingImportEdit {
     /// The edit to add the `TYPE_CHECKING` symbol to the module.
-    type_checking_edit: Edit,
+    type_checking_edit: Option<Edit>,
     /// The edit to add the import to a `TYPE_CHECKING` block.
     add_import_edit: Edit,
 }
 
 impl TypingImportEdit {
-    pub(crate) fn into_edits(self) -> Vec<Edit> {
+    pub(crate) fn into_edits(self) -> (Edit, Option<Edit>) {
-        vec![self.type_checking_edit, self.add_import_edit]
+        if let Some(type_checking_edit) = self.type_checking_edit {
+            (type_checking_edit, Some(self.add_import_edit))
+        } else {
+            (self.add_import_edit, None)
+        }
     }
 }
 

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

@@ -106,6 +106,35 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(
+        Rule::TypingOnlyStandardLibraryImport,
+        Path::new("exempt_type_checking_1.py")
+    )]
+    #[test_case(
+        Rule::TypingOnlyStandardLibraryImport,
+        Path::new("exempt_type_checking_2.py")
+    )]
+    #[test_case(
+        Rule::TypingOnlyStandardLibraryImport,
+        Path::new("exempt_type_checking_3.py")
+    )]
+    fn exempt_type_checking(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!("{}_{}", rule_code.as_ref(), path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("flake8_type_checking").join(path).as_path(),
+            &settings::LinterSettings {
+                flake8_type_checking: super::settings::Settings {
+                    exempt_modules: vec![],
+                    strict: true,
+                    ..Default::default()
+                },
+                ..settings::LinterSettings::for_rule(rule_code)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
+
     #[test_case(
         Rule::RuntimeImportInTypeCheckingBlock,
         Path::new("runtime_evaluated_base_classes_1.py")

crates/ruff_linter/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs

@@ -473,7 +473,9 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) ->
     )?;
 
     // Step 2) Add the import to a `TYPE_CHECKING` block.
-    let add_import_edit = checker.importer().typing_import_edit(
+    let (type_checking_edit, add_import_edit) = checker
+        .importer()
+        .typing_import_edit(
             &ImportedMembers {
                 statement,
                 names: member_names.iter().map(AsRef::as_ref).collect(),
@@ -481,7 +483,8 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) ->
             at,
             checker.semantic(),
             checker.source_type,
-    )?;
+        )?
+        .into_edits();
 
     // Step 3) Quote any runtime usages of the referenced symbol.
     let quote_reference_edits = filter_contained(
@@ -507,10 +510,10 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) ->
     );
 
     Ok(Fix::unsafe_edits(
-        remove_import_edit,
+        type_checking_edit,
         add_import_edit
-            .into_edits()
             .into_iter()
+            .chain(std::iter::once(remove_import_edit))
             .chain(quote_reference_edits),
     )
     .isolate(Checker::isolation(

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-standard-library-import_exempt_type_checking_1.py.snap

@@ -0,0 +1,27 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+exempt_type_checking_1.py:5:20: TCH003 [*] Move standard library import `typing.Final` into a type-checking block
+  |
+3 | from __future__ import annotations
+4 | 
+5 | from typing import Final
+  |                    ^^^^^ TCH003
+6 | 
+7 | Const: Final[dict] = {}
+  |
+  = help: Move into type-checking block
+
+ℹ Unsafe fix
+2 2 | 
+3 3 | from __future__ import annotations
+4 4 | 
+5   |-from typing import Final
+  5 |+from typing import TYPE_CHECKING
+  6 |+
+  7 |+if TYPE_CHECKING:
+  8 |+    from typing import Final
+6 9 | 
+7 10 | Const: Final[dict] = {}
+
+

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-standard-library-import_exempt_type_checking_2.py.snap

@@ -0,0 +1,27 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+exempt_type_checking_2.py:5:20: TCH003 [*] Move standard library import `typing.Final` into a type-checking block
+  |
+3 | from __future__ import annotations
+4 | 
+5 | from typing import Final, TYPE_CHECKING
+  |                    ^^^^^ TCH003
+6 | 
+7 | Const: Final[dict] = {}
+  |
+  = help: Move into type-checking block
+
+ℹ Unsafe fix
+2 2 | 
+3 3 | from __future__ import annotations
+4 4 | 
+5   |-from typing import Final, TYPE_CHECKING
+  5 |+from typing import TYPE_CHECKING
+  6 |+
+  7 |+if TYPE_CHECKING:
+  8 |+    from typing import Final
+6 9 | 
+7 10 | Const: Final[dict] = {}
+
+

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-standard-library-import_exempt_type_checking_3.py.snap

@@ -0,0 +1,28 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+exempt_type_checking_3.py:5:20: TCH003 [*] Move standard library import `typing.Final` into a type-checking block
+  |
+3 | from __future__ import annotations
+4 | 
+5 | from typing import Final, Mapping
+  |                    ^^^^^ TCH003
+6 | 
+7 | Const: Final[dict] = {}
+  |
+  = help: Move into type-checking block
+
+ℹ Unsafe fix
+2  2  | 
+3  3  | from __future__ import annotations
+4  4  | 
+5     |-from typing import Final, Mapping
+   5  |+from typing import Mapping
+   6  |+from typing import TYPE_CHECKING
+   7  |+
+   8  |+if TYPE_CHECKING:
+   9  |+    from typing import Final
+6  10 | 
+7  11 | Const: Final[dict] = {}
+
+

crates/ruff_python_semantic/src/model.rs

@@ -761,6 +761,7 @@ impl<'a> SemanticModel<'a> {
                                     {
                                         return Some(ImportedName {
                                             name: format!("{name}.{member}"),
+                                            source,
                                             range: self.nodes[source].range(),
                                             context: binding.context,
                                         });
@@ -785,6 +786,7 @@ impl<'a> SemanticModel<'a> {
                                         {
                                             return Some(ImportedName {
                                                 name: (*name).to_string(),
+                                                source,
                                                 range: self.nodes[source].range(),
                                                 context: binding.context,
                                             });
@@ -806,6 +808,7 @@ impl<'a> SemanticModel<'a> {
                                     {
                                         return Some(ImportedName {
                                             name: format!("{name}.{member}"),
+                                            source,
                                             range: self.nodes[source].range(),
                                             context: binding.context,
                                         });
@@ -1828,6 +1831,8 @@ pub enum ReadResult {
 pub struct ImportedName {
     /// The name to which the imported symbol is bound.
     name: String,
+    /// The statement from which the symbol is imported.
+    source: NodeId,
     /// The range at which the symbol is imported.
     range: TextRange,
     /// The context in which the symbol is imported.
@@ -1842,6 +1847,10 @@ impl ImportedName {
     pub const fn context(&self) -> ExecutionContext {
         self.context
     }
+
+    pub fn statement<'a>(&self, semantic: &'a SemanticModel) -> &'a Stmt {
+        semantic.statement(self.source)
+    }
 }
 
 impl Ranged for ImportedName {