Support conditional assignment to __lazy_modules__

[ntBre]

Summary

This is a follow-up to #28459 adding simple support for conditionally defining __lazy_modules__. We now track per-module laziness, which allows us to merge definitions like:

__lazy_modules__ = ["foo"]

if condition:
    __lazy_modules__ = ["foo", "bar"]

into known laziness for foo and unknown laziness for bar. As in #28459, that unknown laziness is still helpful and distinct from known eagerness because it suppresses diagnostics.

The one quirk with the branch_id tracking used here is that it doesn't handle exhaustive branches. For example, it would be nice if we noticed that foo is always lazy here:

if condition:
    __lazy_modules__ = ["foo", "bar"]
else:
    __lazy_modules__ = ["foo", "baz"]

but it's instead merged with the implicit known eagerness of all imports before any __lazy_modules__ assignment, producing (foo, Unknown). So this is effectively equivalent to the version without an else for foo.

I think this is a decent tradeoff that biases us against false positives, though, and doesn't involve any extra bookkeeping in if statement traversal.

One of my three Codex reviewers also identified this as a false negative, where the assignment and
use are in the same branch, but we've already marked the module as Unknown:

if condition:
    __lazy_modules__ = ["json"]
    import json
    json.dumps({})

But again we're at least biased in the right direction, and we'd need more bookkeeping to identify
this, so I was okay accepting that for now.

Test Plan

New mdtests for TID254 that exercise this behavior

[astral-sh-bot]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

[MichaReiser]

Codex thinks that supporting if..else isn't too hard? But happy to look at this in a separate PR (or deprioritize)

--- a/crates/ruff_python_semantic/src/model.rs
+++ b/crates/ruff_python_semantic/src/model.rs
@@ -75,6 +75,9 @@
     /// The ID of the current branch.
     branch_id: Option<BranchId>,
 
+    /// The current `if` arm whose direct assignments replace its lazy-module state.
+    lazy_modules_if_arm: Option<BranchId>,
+
     /// Stack of all scopes, along with the identifier of the current scope.
     pub scopes: Scopes<'a>,
     pub scope_id: ScopeId,
@@ -165,7 +168,7 @@
     /// import pathlib  # Lazy.
     /// ```
     ///
-    /// Conditional assignments are merged with the current state:
+    /// Each branch of a module-level `if` starts with the current state and is merged afterward:
     ///
     /// ```python
     /// __lazy_modules__ = ["json"]
@@ -177,9 +180,9 @@
     ///
     /// Here, `json` remains definitely lazy, but `pathlib`'s laziness is unknown.
     ///
-    /// Without an earlier declaration, both modules remain unknown because we merge with the
-    /// default eager state, rather than exhaustively tracking each branch to guarantee an
-    /// assignment occurs.
+    /// Without an earlier declaration, the `if` and `else` states are merged: `json` remains
+    /// lazy, while `pathlib` is unknown. An `if` without an `else` also includes the state before
+    /// the conditional, because its body may not execute.
     pub lazy_modules: Option<LazyModules<'a>>,
 
     /// Exceptions that are handled by the current `try` block.
@@ -223,6 +226,7 @@
             node_id: None,
             branches: Branches::default(),
             branch_id: None,
+            lazy_modules_if_arm: None,
             scopes: Scopes::default(),
             scope_id: ScopeId::global(),
             definitions: Definitions::for_module(module),
@@ -1490,6 +1494,11 @@
         self.branch_id = branch_id;
     }
 
+    /// Let direct assignments replace the lazy-module state within this arm.
+    pub fn set_lazy_modules_if_arm(&mut self, branch_id: Option<BranchId>) {
+        self.lazy_modules_if_arm = branch_id;
+    }
+
     /// Returns an [`Iterator`] over the current statement hierarchy, from the current [`Stmt`]
     /// through to any parents.
     pub fn current_statements(&self) -> impl Iterator<Item = &'a Stmt> + '_ {
@@ -2483,7 +2492,14 @@
             clippy::iter_over_hash_type,
             reason = "each module's laziness is merged independently"
         )]
-        if self.branch_id.is_some() {
+        // Every path reaching the next statement in an arm has executed its direct assignments.
+        // A nested assignment may be skipped, including one inside a loop (which shares the arm's
+        // branch ID), so it must still be merged with the prior state.
+        let direct_if_assignment = self.lazy_modules_if_arm.is_some_and(|arm| {
+            self.branch_id == Some(arm)
+                && matches!(self.current_statements().nth(1), Some(Stmt::If(_)))
+        });
+        if self.branch_id.is_some() && !direct_if_assignment {
             if matches!(self.lazy_modules, Some(LazyModules::Unknown)) {
                 return;
             }
@@ -2528,7 +2544,7 @@
 }
 
 /// Module names and their inferred laziness from `__lazy_modules__` declarations.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum LazyModules<'a> {
     /// Names from literal lists, sets, or tuples, with per-module laziness.
     ///
@@ -2554,6 +2570,41 @@
     Unknown,
 }
 
+impl<'a> LazyModules<'a> {
+    /// Merge the laziness at the exits of two mutually exclusive branches.
+    #[expect(
+        clippy::iter_over_hash_type,
+        reason = "each module's laziness is merged independently"
+    )]
+    pub fn join(left: Option<Self>, right: Option<Self>) -> Option<Self> {
+        match (left, right) {
+            (None, None) => None,
+            (Some(Self::Unknown), _) | (_, Some(Self::Unknown)) => Some(Self::Unknown),
+            (Some(Self::Known(mut left)), Some(Self::Known(right))) => {
+                for (module, laziness) in &mut left {
+                    *laziness = laziness.join(
+                        right
+                            .get(module)
+                            .copied()
+                            .unwrap_or(ImportLaziness::Eager),
+                    );
+                }
+                for (module, laziness) in right {
+                    left.entry(module)
+                        .or_insert_with(|| ImportLaziness::Eager.join(laziness));
+                }
+                Some(Self::Known(left))
+            }
+            (Some(Self::Known(mut modules)), None) | (None, Some(Self::Known(mut modules))) => {
+                for laziness in modules.values_mut() {
+                    *laziness = laziness.join(ImportLaziness::Eager);
+                }
+                Some(Self::Known(modules))
+            }
+        }
+    }
+}
+
 /// Whether an import is lazy, as determined statically.
 ///
 /// Dynamic assignments and conditional membership changes are classified as unknown.
@@ -2565,6 +2616,14 @@
 }
 
 impl ImportLaziness {
+    fn join(self, other: Self) -> Self {
+        if self == other {
+            self
+        } else {
+            Self::Unknown
+        }
+    }
+
     /// Returns `true` if the import laziness is [`Self::Lazy`].
     pub fn is_lazy(&self) -> bool {
         matches!(self, Self::Lazy)
--- a/crates/ruff_linter/src/checkers/ast/mod.rs
+++ b/crates/ruff_linter/src/checkers/ast/mod.rs
@@ -56,8 +56,8 @@
 use ruff_python_semantic::analyze::{imports, typing};
 use ruff_python_semantic::{
     BindingFlags, BindingId, BindingKind, Exceptions, Export, FromImport, GeneratorKind, Globals,
-    Import, ImportLaziness, Module, ModuleKind, ModuleSource, NodeId, ScopeId, ScopeKind,
-    SemanticModel, SemanticModelFlags, StarImport, SubmoduleImport,
+    Import, ImportLaziness, LazyModules, Module, ModuleKind, ModuleSource, NodeId, ScopeId,
+    ScopeKind, SemanticModel, SemanticModelFlags, StarImport, SubmoduleImport,
 };
 use ruff_python_trivia::CommentRanges;
 use ruff_source_file::{OneIndexed, SourceFile, SourceFileBuilder, SourceRow};
@@ -1660,7 +1660,16 @@
             ) => {
                 self.visit_boolean_test(test);
 
-                self.semantic.push_branch();
+                let initial_lazy_modules = self
+                    .semantic
+                    .at_top_level()
+                    .then(|| self.semantic.lazy_modules.clone());
+                let mut merged_lazy_modules = None;
+
+                let branch = self.semantic.push_branch();
+                if initial_lazy_modules.is_some() {
+                    self.semantic.set_lazy_modules_if_arm(branch);
+                }
                 if typing::is_type_checking_block(stmt_if, &self.semantic) {
                     if self.semantic.at_top_level() {
                         self.importer.visit_type_checking_block(stmt);
@@ -1669,13 +1678,42 @@
                 } else {
                     self.visit_body(body);
                 }
+                if let Some(initial) = &initial_lazy_modules {
+                    merged_lazy_modules = std::mem::replace(
+                        &mut self.semantic.lazy_modules,
+                        initial.clone(),
+                    );
+                    self.semantic.set_lazy_modules_if_arm(None);
+                }
                 self.semantic.pop_branch();
 
                 for clause in elif_else_clauses {
-                    self.semantic.push_branch();
+                    let branch = self.semantic.push_branch();
+                    if initial_lazy_modules.is_some() {
+                        self.semantic.set_lazy_modules_if_arm(branch);
+                    }
                     self.visit_elif_else_clause(clause);
+                    if let Some(initial) = &initial_lazy_modules {
+                        merged_lazy_modules = LazyModules::join(
+                            merged_lazy_modules,
+                            std::mem::replace(&mut self.semantic.lazy_modules, initial.clone()),
+                        );
+                        self.semantic.set_lazy_modules_if_arm(None);
+                    }
                     self.semantic.pop_branch();
                 }
+
+                if let Some(initial) = initial_lazy_modules {
+                    // Without an `else`, the body may not execute.
+                    self.semantic.lazy_modules = if elif_else_clauses
+                        .last()
+                        .is_some_and(|clause| clause.test.is_none())
+                    {
+                        merged_lazy_modules
+                    } else {
+                        LazyModules::join(merged_lazy_modules, initial)
+                    };
+                }
             }
             _ => visitor::walk_stmt(self, stmt),
         }
--- a/crates/ruff_linter/resources/mdtest/flake8-tidy-imports/lazy-import-mismatch.md
+++ b/crates/ruff_linter/resources/mdtest/flake8-tidy-imports/lazy-import-mismatch.md
@@ -191,16 +191,79 @@
 
 #### Exhaustive branches
 
-Ruff does not prove that branches cover every path, so `json`'s laziness remains unknown even though
-both branches list it.
+Both branches declare `json`, so it is always lazy. `pathlib` and `math` are declared on
+only one path, so their laziness remains unknown. Modules absent from both branches remain eager.
+
+```py
+if condition:
+    __lazy_modules__ = ["json", "pathlib"]
+    import json  # error: [lazy-import-mismatch]
+else:
+    __lazy_modules__ = ["json", "math"]
+
+import json  # error: [lazy-import-mismatch]
+import pathlib
+import math
+import typing  # error: [lazy-import-mismatch]
+```
+
+#### Exhaustive `elif` branches
+
+An `else` also covers the remaining path after one or more `elif` clauses.
+
+```py
+if first:
+    __lazy_modules__ = ["json"]
+elif second:
+    __lazy_modules__ = ["json"]
+else:
+    __lazy_modules__ = ["json"]
+
+import json  # error: [lazy-import-mismatch]
+```
+
+#### Imports within branches
+
+Each branch starts with the declaration that precedes the conditional. A direct assignment changes
+the laziness of later imports in that branch.
+
+```py
+__lazy_modules__ = ["json"]
+if condition:
+    __lazy_modules__ = []
+    import json
+else:
+    import json  # error: [lazy-import-mismatch]
+```
+
+#### Conditional assignments within branches
+
+The loop in the first branch may not execute, so it cannot make `json` definitely lazy even though
+the other branch declares it.
+
+```py
+if condition:
+    for _ in ():
+        __lazy_modules__ = ["json"]
+else:
+    __lazy_modules__ = ["json"]
+
+import json
+```
+
+#### Multiple declarations within branches
+
+Later declarations in a branch replace earlier ones. Here `json` is eager in one branch and lazy
+in the other, so its laziness after the conditional is unknown.
 
 ```py
 if condition:
     __lazy_modules__ = ["json"]
+    __lazy_modules__ = []
 else:
     __lazy_modules__ = ["json"]
 
-import json  # ok, unknown laziness in our simplified model
+import json

[ntBre]

Sure, not too hard, just doubles the size of the diff 😄 I think I'll merge this as-is for now since I'm still a bit skeptical that conditional assignments in general will be very common, but it's great to have the patch ready if we need it!