Support conditional assignment to __lazy_modules__ (#28491)

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:

```py
__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:

```py
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`:

```py
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

crates/ruff_linter/resources/mdtest/flake8-tidy-imports/lazy-import-mismatch.md

@@ -148,6 +148,61 @@ __lazy_modules__ = configured_modules()
 import json
 ```
 
+### Conditional declarations
+
+```toml
+target-version = "py315"
+
+[lint]
+preview = true
+select = ["TID254"]
+
+[lint.flake8-tidy-imports]
+require-lazy = ["pathlib", "math", "typing"]
+ban-lazy = ["json"]
+```
+
+#### Conditional assignment
+
+The declaration may not execute, so `json`'s laziness is unknown and no diagnostic is emitted.
+
+```py
+if condition:
+    __lazy_modules__ = ["json"]
+
+import json  # ok, may or may not be lazy depending on `condition`
+```
+
+#### Conditional reassignments
+
+Modules present in both declarations remain lazy, and modules absent from both remain eager. Added
+or removed modules have unknown laziness, so their import policies are not enforced.
+
+```py
+__lazy_modules__ = ["json", "pathlib"]
+if condition:
+    __lazy_modules__ = ["json", "math"]
+
+import json  # error: [lazy-import-mismatch]
+import pathlib
+import math
+import typing  # error: [lazy-import-mismatch]
+```
+
+#### Exhaustive branches
+
+Ruff does not prove that branches cover every path, so `json`'s laziness remains unknown even though
+both branches list it.
+
+```py
+if condition:
+    __lazy_modules__ = ["json"]
+else:
+    __lazy_modules__ = ["json"]
+
+import json  # ok, unknown laziness in our simplified model
+```
+
 ### Per-alias policies
 
 Each alias is checked independently. Here `json` must become eager and `pathlib` must become lazy.

crates/ruff_python_semantic/src/model.rs

@@ -155,7 +155,7 @@ pub struct SemanticModel<'a> {
     /// Modules that have been seen by the semantic model.
     pub seen: Modules,
 
-    /// The module names from the most recently visited module-level `__lazy_modules__` assignment.
+    /// Module names and their laziness inferred from module-level `__lazy_modules__` assignments.
     ///
     /// A declaration affects subsequent imports, without changing earlier imports:
     ///
@@ -164,6 +164,22 @@ pub struct SemanticModel<'a> {
     /// __lazy_modules__ = ["json", "pathlib"]
     /// import pathlib  # Lazy.
     /// ```
+    ///
+    /// Conditional assignments are merged with the current state:
+    ///
+    /// ```python
+    /// __lazy_modules__ = ["json"]
+    /// if condition:
+    ///     __lazy_modules__ = ["json", "pathlib"]
+    /// else:
+    ///     __lazy_modules__ = ["json"]
+    /// ```
+    ///
+    /// 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.
     pub lazy_modules: Option<LazyModules<'a>>,
 
     /// Exceptions that are handled by the current `try` block.
@@ -2431,22 +2447,19 @@ impl<'a> SemanticModel<'a> {
     }
 
     /// Test exact module membership in the current `__lazy_modules__` declaration.
-    /// Returns [`ImportLaziness::Unknown`] when the declaration is not a literal collection of strings.
+    /// Returns [`ImportLaziness::Unknown`] for dynamic declarations or uncertain conditional membership.
     pub fn module_laziness(&self, module: &str) -> ImportLaziness {
         match &self.lazy_modules {
             None => ImportLaziness::Eager,
             Some(LazyModules::Unknown) => ImportLaziness::Unknown,
-            Some(LazyModules::Known(modules)) => {
-                if modules.contains(&module) {
-                    ImportLaziness::Lazy
-                } else {
-                    ImportLaziness::Eager
-                }
-            }
+            Some(LazyModules::Known(modules)) => modules
+                .get(module)
+                .copied()
+                .unwrap_or(ImportLaziness::Eager),
         }
     }
 
-    /// Extract literal module names when visiting a `__lazy_modules__` assignment.
+    /// Extract literal module names, merging conditional assignments with the current state.
     pub fn set_lazy_modules(&mut self, value: &'a Expr) {
         let names = match value {
             Expr::List(ast::ExprList { elts, .. })
@@ -2456,12 +2469,41 @@ impl<'a> SemanticModel<'a> {
                 .map(|element| {
                     element
                         .as_string_literal_expr()
-                        .map(|literal| literal.value.to_str())
+                        .map(|literal| (literal.value.to_str(), ImportLaziness::Lazy))
                 })
-                .collect::<Option<Vec<_>>>(),
+                .collect::<Option<FxHashMap<_, _>>>(),
             _ => None,
         };
-        self.lazy_modules = Some(names.map_or(LazyModules::Unknown, LazyModules::Known));
+        let Some(mut modules) = names else {
+            self.lazy_modules = Some(LazyModules::Unknown);
+            return;
+        };
+
+        #[expect(
+            clippy::iter_over_hash_type,
+            reason = "each module's laziness is merged independently"
+        )]
+        if self.branch_id.is_some() {
+            if matches!(self.lazy_modules, Some(LazyModules::Unknown)) {
+                return;
+            }
+
+            // The assignment may not execute. Only modules already known to be lazy remain so.
+            for (module, laziness) in &mut modules {
+                if !self.module_laziness(module).is_lazy() {
+                    *laziness = ImportLaziness::Unknown;
+                }
+            }
+
+            // A removed entry may still be lazy on paths where the assignment does not execute.
+            if let Some(LazyModules::Known(previous)) = &self.lazy_modules {
+                for module in previous.keys() {
+                    modules.entry(module).or_insert(ImportLaziness::Unknown);
+                }
+            }
+        }
+
+        self.lazy_modules = Some(LazyModules::Known(modules));
     }
 
     /// Return the module tested for membership in `__lazy_modules__`.
@@ -2485,17 +2527,19 @@ impl<'a> SemanticModel<'a> {
     }
 }
 
-/// Statically known module names in a `__lazy_modules__` declaration.
+/// Module names and their inferred laziness from `__lazy_modules__` declarations.
 #[derive(Debug)]
 pub enum LazyModules<'a> {
-    /// A literal list, set, or tuple that can be analyzed.
+    /// Names from literal lists, sets, or tuples, with per-module laziness.
     ///
     /// For example:
     ///
     /// ```py
     /// __lazy_modules__ = ["a", "list"]
     /// ```
-    Known(Vec<&'a str>),
+    ///
+    /// Conditional assignments can leave a listed name's laziness unknown.
+    Known(FxHashMap<&'a str, ImportLaziness>),
 
     /// The declaration is present but not a literal collection that can be analyzed.
     ///
@@ -2512,7 +2556,7 @@ pub enum LazyModules<'a> {
 
 /// Whether an import is lazy, as determined statically.
 ///
-/// Dynamic assignments to `__lazy_modules__` are classified as unknown.
+/// Dynamic assignments and conditional membership changes are classified as unknown.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum ImportLaziness {
     Lazy,