Avoid if-else simplification for TYPE_CHECKING blocks

Author: charliermarsh

crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM108.py

@@ -126,3 +126,12 @@ def f():
         x = yield 3
     else:
         x = yield 5
+
+
+from typing import TYPE_CHECKING
+
+# OK
+if TYPE_CHECKING:
+    x = 3
+else:
+    x = 5

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1044,16 +1044,16 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 flake8_simplify::rules::if_with_same_arms(checker, checker.locator, if_);
             }
             if checker.enabled(Rule::NeedlessBool) {
-                flake8_simplify::rules::needless_bool(checker, stmt);
+                flake8_simplify::rules::needless_bool(checker, if_);
             }
             if checker.enabled(Rule::IfElseBlockInsteadOfDictLookup) {
-                flake8_simplify::rules::manual_dict_lookup(checker, if_);
+                flake8_simplify::rules::if_else_block_instead_of_dict_lookup(checker, if_);
             }
             if checker.enabled(Rule::IfElseBlockInsteadOfIfExp) {
-                flake8_simplify::rules::use_ternary_operator(checker, stmt);
+                flake8_simplify::rules::if_else_block_instead_of_if_exp(checker, if_);
             }
             if checker.enabled(Rule::IfElseBlockInsteadOfDictGet) {
-                flake8_simplify::rules::use_dict_get_with_default(checker, if_);
+                flake8_simplify::rules::if_else_block_instead_of_dict_get(checker, if_);
             }
             if checker.enabled(Rule::TypeCheckWithoutTypeError) {
                 tryceratops::rules::type_check_without_type_error(

crates/ruff_linter/src/rules/flake8_simplify/rules/collapsible_if.rs

@@ -9,6 +9,7 @@ use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::{self as ast, whitespace, Constant, ElifElseClause, Expr, Stmt};
 use ruff_python_codegen::Stylist;
+use ruff_python_semantic::analyze::typing::{is_sys_version_block, is_type_checking_block};
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextRange};
@@ -96,6 +97,16 @@ pub(crate) fn nested_if_statements(
         return;
     };
 
+    // Avoid suggesting ternary for `if sys.version_info >= ...`-style checks.
+    if is_sys_version_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
+    // Avoid suggesting ternary for `if TYPE_CHECKING:`-style checks.
+    if is_type_checking_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
     let mut diagnostic = Diagnostic::new(
         CollapsibleIf,
         TextRange::new(nested_if.start(), colon.end()),

crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_get.rs

@@ -5,6 +5,7 @@ use ruff_python_ast::helpers::contains_effect;
 use ruff_python_ast::{
     self as ast, Arguments, CmpOp, ElifElseClause, Expr, ExprContext, Identifier, Stmt,
 };
+use ruff_python_semantic::analyze::typing::{is_sys_version_block, is_type_checking_block};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
@@ -55,7 +56,7 @@ impl Violation for IfElseBlockInsteadOfDictGet {
 }
 
 /// SIM401
-pub(crate) fn use_dict_get_with_default(checker: &mut Checker, stmt_if: &ast::StmtIf) {
+pub(crate) fn if_else_block_instead_of_dict_get(checker: &mut Checker, stmt_if: &ast::StmtIf) {
     let ast::StmtIf {
         test,
         body,
@@ -136,6 +137,16 @@ pub(crate) fn use_dict_get_with_default(checker: &mut Checker, stmt_if: &ast::St
         return;
     }
 
+    // Avoid suggesting ternary for `if sys.version_info >= ...`-style checks.
+    if is_sys_version_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
+    // Avoid suggesting ternary for `if TYPE_CHECKING:`-style checks.
+    if is_type_checking_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
     // Check that the default value is not "complex".
     if contains_effect(default_value, |id| checker.semantic().is_builtin(id)) {
         return;

crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_lookup.rs

@@ -5,6 +5,7 @@ use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::comparable::ComparableConstant;
 use ruff_python_ast::helpers::contains_effect;
 use ruff_python_ast::{self as ast, CmpOp, ElifElseClause, Expr, Stmt};
+use ruff_python_semantic::analyze::typing::{is_sys_version_block, is_type_checking_block};
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -39,7 +40,7 @@ impl Violation for IfElseBlockInsteadOfDictLookup {
     }
 }
 /// SIM116
-pub(crate) fn manual_dict_lookup(checker: &mut Checker, stmt_if: &ast::StmtIf) {
+pub(crate) fn if_else_block_instead_of_dict_lookup(checker: &mut Checker, stmt_if: &ast::StmtIf) {
     // Throughout this rule:
     // * Each if or elif statement's test must consist of a constant equality check with the same variable.
     // * Each if or elif statement's body must consist of a single `return`.
@@ -75,13 +76,24 @@ pub(crate) fn manual_dict_lookup(checker: &mut Checker, stmt_if: &ast::StmtIf) {
     let [Stmt::Return(ast::StmtReturn { value, range: _ })] = body.as_slice() else {
         return;
     };
+
     if value
         .as_ref()
         .is_some_and(|value| contains_effect(value, |id| checker.semantic().is_builtin(id)))
     {
         return;
     }
 
+    // Avoid suggesting ternary for `if sys.version_info >= ...`-style checks.
+    if is_sys_version_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
+    // Avoid suggesting ternary for `if TYPE_CHECKING:`-style checks.
+    if is_type_checking_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
     let mut constants: FxHashSet<ComparableConstant> = FxHashSet::default();
     constants.insert(constant.into());
 

crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_if_exp.rs

@@ -1,8 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::any_over_expr;
 use ruff_python_ast::{self as ast, ElifElseClause, Expr, Stmt};
-use ruff_python_semantic::SemanticModel;
+use ruff_python_semantic::analyze::typing::{is_sys_version_block, is_type_checking_block};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
@@ -51,16 +50,14 @@ impl Violation for IfElseBlockInsteadOfIfExp {
 }
 
 /// SIM108
-pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) {
-    let Stmt::If(ast::StmtIf {
+pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &ast::StmtIf) {
+    let ast::StmtIf {
         test,
         body,
         elif_else_clauses,
         range: _,
-    }) = stmt
-    else {
-        return;
-    };
+    } = stmt_if;
+
     // `test: None` to only match an `else` clause
     let [ElifElseClause {
         body: else_body,
@@ -99,13 +96,6 @@ pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) {
         return;
     }
 
-    // Avoid suggesting ternary for `if sys.version_info >= ...`-style and
-    // `if sys.platform.startswith("...")`-style checks.
-    let ignored_call_paths: &[&[&str]] = &[&["sys", "version_info"], &["sys", "platform"]];
-    if contains_call_path(test, ignored_call_paths, checker.semantic()) {
-        return;
-    }
-
     // Avoid suggesting ternary for `if (yield ...)`-style checks.
     // TODO(charlie): Fix precedence handling for yields in generator.
     if matches!(
@@ -121,14 +111,24 @@ pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) {
         return;
     }
 
+    // Avoid suggesting ternary for `if sys.version_info >= ...`-style checks.
+    if is_sys_version_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
+    // Avoid suggesting ternary for `if TYPE_CHECKING:`-style checks.
+    if is_type_checking_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
     let target_var = &body_target;
     let ternary = ternary(target_var, body_value, test, else_value);
     let contents = checker.generator().stmt(&ternary);
 
     // Don't flag if the resulting expression would exceed the maximum line length.
     if !fits(
         &contents,
-        stmt.into(),
+        stmt_if.into(),
         checker.locator(),
         checker.settings.line_length,
         checker.settings.tab_size,
@@ -140,26 +140,17 @@ pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) {
         IfElseBlockInsteadOfIfExp {
             contents: contents.clone(),
         },
-        stmt.range(),
+        stmt_if.range(),
     );
-    if !checker.indexer().has_comments(stmt, checker.locator()) {
+    if !checker.indexer().has_comments(stmt_if, checker.locator()) {
         diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
             contents,
-            stmt.range(),
+            stmt_if.range(),
         )));
     }
     checker.diagnostics.push(diagnostic);
 }
 
-/// Return `true` if the `Expr` contains a reference to any of the given `${module}.${target}`.
-fn contains_call_path(expr: &Expr, targets: &[&[&str]], semantic: &SemanticModel) -> bool {
-    any_over_expr(expr, &|expr| {
-        semantic
-            .resolve_call_path(expr)
-            .is_some_and(|call_path| targets.iter().any(|target| &call_path.as_slice() == target))
-    })
-}
-
 fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Expr) -> Stmt {
     let node = ast::ExprIfExp {
         test: Box::new(test.clone()),

crates/ruff_linter/src/rules/flake8_simplify/rules/needless_bool.rs

@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast, Arguments, Constant, ElifElseClause, Expr, ExprContext, Stmt};
+use ruff_python_semantic::analyze::typing::{is_sys_version_block, is_type_checking_block};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
@@ -48,24 +49,22 @@ impl Violation for NeedlessBool {
 }
 
 /// SIM103
-pub(crate) fn needless_bool(checker: &mut Checker, stmt: &Stmt) {
-    let Stmt::If(ast::StmtIf {
+pub(crate) fn needless_bool(checker: &mut Checker, stmt_if: &ast::StmtIf) {
+    let ast::StmtIf {
         test: if_test,
         body: if_body,
         elif_else_clauses,
         range: _,
-    }) = stmt
-    else {
-        return;
-    };
+    } = stmt_if;
+
     // Extract an `if` or `elif` (that returns) followed by an else (that returns the same value)
     let (if_test, if_body, else_body, range) = match elif_else_clauses.as_slice() {
         // if-else case
         [ElifElseClause {
             body: else_body,
             test: None,
             ..
-        }] => (if_test.as_ref(), if_body, else_body, stmt.range()),
+        }] => (if_test.as_ref(), if_body, else_body, stmt_if.range()),
         // elif-else case
         [.., ElifElseClause {
             body: elif_body,
@@ -97,6 +96,16 @@ pub(crate) fn needless_bool(checker: &mut Checker, stmt: &Stmt) {
         return;
     }
 
+    // Avoid suggesting ternary for `if sys.version_info >= ...`-style checks.
+    if is_sys_version_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
+    // Avoid suggesting ternary for `if TYPE_CHECKING:`-style checks.
+    if is_type_checking_block(stmt_if, checker.semantic()) {
+        return;
+    }
+
     let condition = checker.generator().expr(if_test);
     let mut diagnostic = Diagnostic::new(NeedlessBool { condition }, range);
     if matches!(if_return, Bool::True)

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -1,7 +1,7 @@
 //! Analysis rules for the `typing` module.
 
 use ruff_python_ast::call_path::{from_qualified_name, from_unqualified_name, CallPath};
-use ruff_python_ast::helpers::{is_const_false, map_subscript};
+use ruff_python_ast::helpers::{any_over_expr, is_const_false, map_subscript};
 use ruff_python_ast::{
     self as ast, Constant, Expr, Int, Operator, ParameterWithDefault, Parameters, Stmt,
 };
@@ -302,7 +302,7 @@ pub fn is_mutable_expr(expr: &Expr, semantic: &SemanticModel) -> bool {
     }
 }
 
-/// Return `true` if [`Expr`] is a guard for a type-checking block.
+/// Return `true` if [`ast::StmtIf`] is a guard for a type-checking block.
 pub fn is_type_checking_block(stmt: &ast::StmtIf, semantic: &SemanticModel) -> bool {
     let ast::StmtIf { test, .. } = stmt;
 
@@ -333,6 +333,17 @@ pub fn is_type_checking_block(stmt: &ast::StmtIf, semantic: &SemanticModel) -> b
     false
 }
 
+/// Returns `true` if the [`ast::StmtIf`] is a version-checking block (e.g., `if sys.version_info >= ...:`).
+pub fn is_sys_version_block(stmt: &ast::StmtIf, semantic: &SemanticModel) -> bool {
+    let ast::StmtIf { test, .. } = stmt;
+
+    any_over_expr(test, &|expr| {
+        semantic.resolve_call_path(expr).is_some_and(|call_path| {
+            matches!(call_path.as_slice(), ["sys", "version_info" | "platform"])
+        })
+    })
+}
+
 /// Abstraction for a type checker, conservatively checks for the intended type(s).
 trait TypeChecker {
     /// Check annotation expression to match the intended type(s).