Track TYPE_CHECKING blocks in Importer (#4593)

Cargo.lock

@@ -2012,6 +2012,7 @@ dependencies = [
  "bitflags 2.3.1",
  "is-macro",
  "nohash-hasher",
+ "num-traits",
  "ruff_index",
  "ruff_python_ast",
  "ruff_python_stdlib",

Cargo.toml

@@ -26,6 +26,8 @@ libcst = { git = "https://github.com/charliermarsh/LibCST", rev = "80e4c1399f95e
 log = { version = "0.4.17" }
 memchr = "2.5.0"
 nohash-hasher = { version = "0.2.0" }
+num-bigint = { version = "0.4.3" }
+num-traits = { version = "0.2.15" }
 once_cell = { version = "1.17.1" }
 path-absolutize = { version = "3.0.14" }
 proc-macro2 = { version = "1.0.51" }

crates/ruff/Cargo.toml

@@ -41,8 +41,8 @@ libcst = { workspace = true }
 log = { workspace = true }
 natord = { version = "1.0.9" }
 nohash-hasher = { workspace = true }
-num-bigint = { version = "0.4.3" }
+num-bigint = { workspace = true }
-num-traits = { version = "0.2.15" }
+num-traits = { workspace = true }
 once_cell = { workspace = true }
 path-absolutize = { workspace = true, features = [
   "once_cell_cache",

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

@@ -2076,16 +2076,21 @@ where
                 self.visit_body(body);
                 self.visit_body(orelse);
             }
-            Stmt::If(ast::StmtIf {
+            Stmt::If(
-                test,
+                stmt_if @ ast::StmtIf {
-                body,
+                    test,
-                orelse,
+                    body,
-                range: _,
+                    orelse,
-            }) => {
+                    range: _,
+                },
+            ) => {
                 self.visit_boolean_test(test);
 
-                if flake8_type_checking::helpers::is_type_checking_block(&self.semantic_model, test)
+                if analyze::typing::is_type_checking_block(stmt_if, &self.semantic_model) {
-                {
+                    if self.semantic_model.at_top_level() {
+                        self.importer.visit_type_checking_block(stmt);
+                    }
+
                     if self.enabled(Rule::EmptyTypeCheckingBlock) {
                         flake8_type_checking::rules::empty_type_checking_block(self, stmt, body);
                     }

crates/ruff/src/importer/mod.rs

@@ -18,10 +18,16 @@ use crate::importer::insertion::Insertion;
 mod insertion;
 
 pub(crate) struct Importer<'a> {
+    /// The Python AST to which we are adding imports.
     python_ast: &'a Suite,
+    /// The [`Locator`] for the Python AST.
     locator: &'a Locator<'a>,
+    /// The [`Stylist`] for the Python AST.
     stylist: &'a Stylist<'a>,
-    ordered_imports: Vec<&'a Stmt>,
+    /// The list of visited, top-level runtime imports in the Python AST.
+    runtime_imports: Vec<&'a Stmt>,
+    /// The list of visited, top-level `if TYPE_CHECKING:` blocks in the Python AST.
+    type_checking_blocks: Vec<&'a Stmt>,
 }
 
 impl<'a> Importer<'a> {
@@ -34,13 +40,19 @@ impl<'a> Importer<'a> {
             python_ast,
             locator,
             stylist,
-            ordered_imports: Vec::default(),
+            runtime_imports: Vec::default(),
+            type_checking_blocks: Vec::default(),
         }
     }
 
     /// Visit a top-level import statement.
     pub(crate) fn visit_import(&mut self, import: &'a Stmt) {
-        self.ordered_imports.push(import);
+        self.runtime_imports.push(import);
+    }
+
+    /// Visit a top-level type-checking block.
+    pub(crate) fn visit_type_checking_block(&mut self, type_checking_block: &'a Stmt) {
+        self.type_checking_blocks.push(type_checking_block);
     }
 
     /// Add an import statement to import the given module.
@@ -168,17 +180,17 @@ impl<'a> Importer<'a> {
 
     /// Return the import statement that precedes the given position, if any.
     fn preceding_import(&self, at: TextSize) -> Option<&Stmt> {
-        self.ordered_imports
+        self.runtime_imports
             .partition_point(|stmt| stmt.start() < at)
             .checked_sub(1)
-            .map(|idx| self.ordered_imports[idx])
+            .map(|idx| self.runtime_imports[idx])
     }
 
     /// Return the top-level [`Stmt`] that imports the given module using `Stmt::ImportFrom`
     /// preceding the given position, if any.
     fn find_import_from(&self, module: &str, at: TextSize) -> Option<&Stmt> {
         let mut import_from = None;
-        for stmt in &self.ordered_imports {
+        for stmt in &self.runtime_imports {
             if stmt.start() >= at {
                 break;
             }

crates/ruff/src/rules/flake8_type_checking/helpers.rs

@@ -1,48 +1,9 @@
-use num_traits::Zero;
-use rustpython_parser::ast::{self, Constant, Expr};
-
 use ruff_python_ast::call_path::from_qualified_name;
 use ruff_python_ast::helpers::map_callable;
 use ruff_python_semantic::binding::{Binding, BindingKind};
 use ruff_python_semantic::model::SemanticModel;
 use ruff_python_semantic::scope::ScopeKind;
-
+use rustpython_parser::ast;
-/// Return `true` if [`Expr`] is a guard for a type-checking block.
-pub(crate) fn is_type_checking_block(semantic_model: &SemanticModel, test: &Expr) -> bool {
-    // Ex) `if False:`
-    if matches!(
-        test,
-        Expr::Constant(ast::ExprConstant {
-            value: Constant::Bool(false),
-            ..
-        })
-    ) {
-        return true;
-    }
-
-    // Ex) `if 0:`
-    if let Expr::Constant(ast::ExprConstant {
-        value: Constant::Int(value),
-        ..
-    }) = &test
-    {
-        if value.is_zero() {
-            return true;
-        }
-    }
-
-    // Ex) `if typing.TYPE_CHECKING:`
-    if semantic_model
-        .resolve_call_path(test)
-        .map_or(false, |call_path| {
-            call_path.as_slice() == ["typing", "TYPE_CHECKING"]
-        })
-    {
-        return true;
-    }
-
-    false
-}
 
 pub(crate) fn is_valid_runtime_import(semantic_model: &SemanticModel, binding: &Binding) -> bool {
     if matches!(

crates/ruff_python_ast/Cargo.toml

@@ -16,8 +16,8 @@ is-macro = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 memchr = { workspace = true }
-num-bigint = { version = "0.4.3" }
+num-bigint = { workspace = true }
-num-traits = { version = "0.2.15" }
+num-traits = { workspace = true }
 once_cell = { workspace = true }
 rustc-hash = { workspace = true }
 rustpython-literal = { workspace = true }

crates/ruff_python_semantic/Cargo.toml

@@ -16,6 +16,7 @@ ruff_index = { path = "../ruff_index" }
 bitflags = { workspace = true }
 is-macro = { workspace = true }
 nohash-hasher = { workspace = true }
+num-traits = { workspace = true }
 rustc-hash = { workspace = true }
 rustpython-parser = { workspace = true }
 smallvec = { workspace = true }

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -1,5 +1,6 @@
 use rustpython_parser::ast::{self, Constant, Expr, Operator};
 
+use num_traits::identities::Zero;
 use ruff_python_ast::call_path::{from_unqualified_name, CallPath};
 use ruff_python_stdlib::typing::{
     IMMUTABLE_GENERIC_TYPES, IMMUTABLE_TYPES, PEP_585_GENERICS, PEP_593_SUBSCRIPTS, SUBSCRIPTS,
@@ -26,39 +27,41 @@ pub enum SubscriptKind {
 
 pub fn match_annotated_subscript<'a>(
     expr: &Expr,
-    model: &SemanticModel,
+    semantic_model: &SemanticModel,
     typing_modules: impl Iterator<Item = &'a str>,
 ) -> Option<SubscriptKind> {
     if !matches!(expr, Expr::Name(_) | Expr::Attribute(_)) {
         return None;
     }
 
-    model.resolve_call_path(expr).and_then(|call_path| {
+    semantic_model
-        if SUBSCRIPTS.contains(&call_path.as_slice()) {
+        .resolve_call_path(expr)
-            return Some(SubscriptKind::AnnotatedSubscript);
+        .and_then(|call_path| {
-        }
+            if SUBSCRIPTS.contains(&call_path.as_slice()) {
-        if PEP_593_SUBSCRIPTS.contains(&call_path.as_slice()) {
+                return Some(SubscriptKind::AnnotatedSubscript);
-            return Some(SubscriptKind::PEP593AnnotatedSubscript);
+            }
-        }
+            if PEP_593_SUBSCRIPTS.contains(&call_path.as_slice()) {
+                return Some(SubscriptKind::PEP593AnnotatedSubscript);
+            }
 
-        for module in typing_modules {
+            for module in typing_modules {
-            let module_call_path: CallPath = from_unqualified_name(module);
+                let module_call_path: CallPath = from_unqualified_name(module);
-            if call_path.starts_with(&module_call_path) {
+                if call_path.starts_with(&module_call_path) {
-                for subscript in SUBSCRIPTS.iter() {
+                    for subscript in SUBSCRIPTS.iter() {
-                    if call_path.last() == subscript.last() {
+                        if call_path.last() == subscript.last() {
-                        return Some(SubscriptKind::AnnotatedSubscript);
+                            return Some(SubscriptKind::AnnotatedSubscript);
+                        }
                     }
-                }
+                    for subscript in PEP_593_SUBSCRIPTS.iter() {
-                for subscript in PEP_593_SUBSCRIPTS.iter() {
+                        if call_path.last() == subscript.last() {
-                    if call_path.last() == subscript.last() {
+                            return Some(SubscriptKind::PEP593AnnotatedSubscript);
-                        return Some(SubscriptKind::PEP593AnnotatedSubscript);
+                        }
                     }
                 }
             }
-        }
 
-        None
+            None
-    })
+        })
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
@@ -80,25 +83,27 @@ impl std::fmt::Display for ModuleMember {
 
 /// Returns the PEP 585 standard library generic variant for a `typing` module reference, if such
 /// a variant exists.
-pub fn to_pep585_generic(expr: &Expr, model: &SemanticModel) -> Option<ModuleMember> {
+pub fn to_pep585_generic(expr: &Expr, semantic_model: &SemanticModel) -> Option<ModuleMember> {
-    model.resolve_call_path(expr).and_then(|call_path| {
+    semantic_model
-        let [module, name] = call_path.as_slice() else {
+        .resolve_call_path(expr)
+        .and_then(|call_path| {
+            let [module, name] = call_path.as_slice() else {
             return None;
         };
-        PEP_585_GENERICS
+            PEP_585_GENERICS.iter().find_map(
-            .iter()
+                |((from_module, from_member), (to_module, to_member))| {
-            .find_map(|((from_module, from_member), (to_module, to_member))| {
+                    if module == from_module && name == from_member {
-                if module == from_module && name == from_member {
+                        if to_module.is_empty() {
-                    if to_module.is_empty() {
+                            Some(ModuleMember::BuiltIn(to_member))
-                        Some(ModuleMember::BuiltIn(to_member))
+                        } else {
+                            Some(ModuleMember::Member(to_module, to_member))
+                        }
                     } else {
-                        Some(ModuleMember::Member(to_module, to_member))
+                        None
                     }
-                } else {
+                },
-                    None
+            )
-                }
+        })
-            })
-    })
 }
 
 /// Return whether a given expression uses a PEP 585 standard library generic.
@@ -128,7 +133,7 @@ pub enum Pep604Operator {
 pub fn to_pep604_operator(
     value: &Expr,
     slice: &Expr,
-    model: &SemanticModel,
+    semantic_model: &SemanticModel,
 ) -> Option<Pep604Operator> {
     /// Returns `true` if any argument in the slice is a string.
     fn any_arg_is_str(slice: &Expr) -> bool {
@@ -148,13 +153,13 @@ pub fn to_pep604_operator(
         return None;
     }
 
-    model
+    semantic_model
         .resolve_call_path(value)
         .as_ref()
         .and_then(|call_path| {
-            if model.match_typing_call_path(call_path, "Optional") {
+            if semantic_model.match_typing_call_path(call_path, "Optional") {
                 Some(Pep604Operator::Optional)
-            } else if model.match_typing_call_path(call_path, "Union") {
+            } else if semantic_model.match_typing_call_path(call_path, "Union") {
                 Some(Pep604Operator::Union)
             } else {
                 None
@@ -164,18 +169,21 @@ pub fn to_pep604_operator(
 
 /// Return `true` if `Expr` represents a reference to a type annotation that resolves to an
 /// immutable type.
-pub fn is_immutable_annotation(model: &SemanticModel, expr: &Expr) -> bool {
+pub fn is_immutable_annotation(semantic_model: &SemanticModel, expr: &Expr) -> bool {
     match expr {
         Expr::Name(_) | Expr::Attribute(_) => {
-            model.resolve_call_path(expr).map_or(false, |call_path| {
+            semantic_model
-                IMMUTABLE_TYPES
+                .resolve_call_path(expr)
-                    .iter()
+                .map_or(false, |call_path| {
-                    .chain(IMMUTABLE_GENERIC_TYPES)
+                    IMMUTABLE_TYPES
-                    .any(|target| call_path.as_slice() == *target)
+                        .iter()
-            })
+                        .chain(IMMUTABLE_GENERIC_TYPES)
+                        .any(|target| call_path.as_slice() == *target)
+                })
         }
-        Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
+        Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => semantic_model
-            model.resolve_call_path(value).map_or(false, |call_path| {
+            .resolve_call_path(value)
+            .map_or(false, |call_path| {
                 if IMMUTABLE_GENERIC_TYPES
                     .iter()
                     .any(|target| call_path.as_slice() == *target)
@@ -183,30 +191,33 @@ pub fn is_immutable_annotation(model: &SemanticModel, expr: &Expr) -> bool {
                     true
                 } else if call_path.as_slice() == ["typing", "Union"] {
                     if let Expr::Tuple(ast::ExprTuple { elts, .. }) = slice.as_ref() {
-                        elts.iter().all(|elt| is_immutable_annotation(model, elt))
+                        elts.iter()
+                            .all(|elt| is_immutable_annotation(semantic_model, elt))
                     } else {
                         false
                     }
                 } else if call_path.as_slice() == ["typing", "Optional"] {
-                    is_immutable_annotation(model, slice)
+                    is_immutable_annotation(semantic_model, slice)
                 } else if call_path.as_slice() == ["typing", "Annotated"] {
                     if let Expr::Tuple(ast::ExprTuple { elts, .. }) = slice.as_ref() {
                         elts.first()
-                            .map_or(false, |elt| is_immutable_annotation(model, elt))
+                            .map_or(false, |elt| is_immutable_annotation(semantic_model, elt))
                     } else {
                         false
                     }
                 } else {
                     false
                 }
-            })
+            }),
-        }
         Expr::BinOp(ast::ExprBinOp {
             left,
             op: Operator::BitOr,
             right,
             range: _range,
-        }) => is_immutable_annotation(model, left) && is_immutable_annotation(model, right),
+        }) => {
+            is_immutable_annotation(semantic_model, left)
+                && is_immutable_annotation(semantic_model, right)
+        }
         Expr::Constant(ast::ExprConstant {
             value: Constant::None,
             ..
@@ -238,16 +249,57 @@ const IMMUTABLE_FUNCS: &[&[&str]] = &[
 
 /// Return `true` if `func` is a function that returns an immutable object.
 pub fn is_immutable_func(
-    model: &SemanticModel,
+    semantic_model: &SemanticModel,
     func: &Expr,
     extend_immutable_calls: &[CallPath],
 ) -> bool {
-    model.resolve_call_path(func).map_or(false, |call_path| {
+    semantic_model
-        IMMUTABLE_FUNCS
+        .resolve_call_path(func)
-            .iter()
+        .map_or(false, |call_path| {
-            .any(|target| call_path.as_slice() == *target)
+            IMMUTABLE_FUNCS
-            || extend_immutable_calls
                 .iter()
-                .any(|target| call_path == *target)
+                .any(|target| call_path.as_slice() == *target)
-    })
+                || extend_immutable_calls
+                    .iter()
+                    .any(|target| call_path == *target)
+        })
+}
+
+/// Return `true` if [`Expr`] is a guard for a type-checking block.
+pub fn is_type_checking_block(stmt: &ast::StmtIf, semantic_model: &SemanticModel) -> bool {
+    let ast::StmtIf { test, .. } = stmt;
+
+    // Ex) `if False:`
+    if matches!(
+        test.as_ref(),
+        Expr::Constant(ast::ExprConstant {
+            value: Constant::Bool(false),
+            ..
+        })
+    ) {
+        return true;
+    }
+
+    // Ex) `if 0:`
+    if let Expr::Constant(ast::ExprConstant {
+        value: Constant::Int(value),
+        ..
+    }) = test.as_ref()
+    {
+        if value.is_zero() {
+            return true;
+        }
+    }
+
+    // Ex) `if typing.TYPE_CHECKING:`
+    if semantic_model
+        .resolve_call_path(test)
+        .map_or(false, |call_path| {
+            call_path.as_slice() == ["typing", "TYPE_CHECKING"]
+        })
+    {
+        return true;
+    }
+
+    false
 }