[pyflakes] Revert to stable behavior if imports for module lie in alternate branches for F401 (#20878)

Closes #20839

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

@@ -528,6 +528,38 @@ mod tests {
         import a",
         "f401_use_in_between_imports"
     )]
+    #[test_case(
+        r"
+        if cond:
+            import a
+            import a.b
+            a.foo()
+        ",
+        "f401_same_branch"
+    )]
+    #[test_case(
+        r"
+        try:
+            import a.b.c
+        except ImportError:
+            import argparse
+            import a
+            a.b = argparse.Namespace()
+        ",
+        "f401_different_branch"
+    )]
+    #[test_case(
+        r"
+        import mlflow.pyfunc.loaders.chat_agent
+        import mlflow.pyfunc.loaders.chat_model
+        import mlflow.pyfunc.loaders.code_model
+        from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+        if IS_PYDANTIC_V2_OR_NEWER:
+            import mlflow.pyfunc.loaders.responses_agent
+        ",
+        "f401_type_checking"
+    )]
     fn f401_preview_refined_submodule_handling(contents: &str, snapshot: &str) {
         let diagnostics = test_contents(
             &SourceKind::Python(dedent(contents).to_string()),

crates/ruff_linter/src/rules/pyflakes/rules/unused_import.rs

@@ -898,6 +898,10 @@ fn best_match<'a, 'b>(
 
 #[inline]
 fn has_simple_shadowed_bindings(scope: &Scope, id: BindingId, semantic: &SemanticModel) -> bool {
+    let Some(binding_node) = semantic.binding(id).source else {
+        return false;
+    };
+
     scope.shadowed_bindings(id).enumerate().all(|(i, shadow)| {
         let shadowed_binding = semantic.binding(shadow);
         // Bail if one of the shadowed bindings is
@@ -912,6 +916,34 @@ fn has_simple_shadowed_bindings(scope: &Scope, id: BindingId, semantic: &Semanti
         if i > 0 && shadowed_binding.is_used() {
             return false;
         }
+        // We want to allow a situation like this:
+        //
+        // ```python
+        // import a.b
+        // if TYPE_CHECKING:
+        //     import a.b.c
+        // ```
+        // but bail in a situation like this:
+        //
+        // ```python
+        // try:
+        //     import a.b
+        // except ImportError:
+        //     import argparse
+        //     import a
+        //     a.b = argparse.Namespace()
+        // ```
+        //
+        // So we require that all the shadowed bindings dominate the
+        // last live binding for the import. That is: if the last live
+        // binding is executed it should imply that all the shadowed
+        // bindings were executed as well.
+        if shadowed_binding
+            .source
+            .is_none_or(|node_id| !semantic.dominates(node_id, binding_node))
+        {
+            return false;
+        }
         matches!(
             shadowed_binding.kind,
             BindingKind::Import(_) | BindingKind::SubmoduleImport(_)

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__f401_different_branch.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/pyflakes/mod.rs
+---
+

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__f401_same_branch.snap

@@ -0,0 +1,18 @@
+---
+source: crates/ruff_linter/src/rules/pyflakes/mod.rs
+---
+F401 [*] `a.b` imported but unused
+ --> f401_preview_submodule.py:4:12
+  |
+2 | if cond:
+3 |     import a
+4 |     import a.b
+  |            ^^^
+5 |     a.foo()
+  |
+help: Remove unused import: `a.b`
+1 | 
+2 | if cond:
+3 |     import a
+  -     import a.b
+4 |     a.foo()

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__f401_type_checking.snap

@@ -0,0 +1,66 @@
+---
+source: crates/ruff_linter/src/rules/pyflakes/mod.rs
+---
+F401 [*] `mlflow.pyfunc.loaders.chat_agent` imported but unused
+ --> f401_preview_submodule.py:2:8
+  |
+2 | import mlflow.pyfunc.loaders.chat_agent
+  |        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+3 | import mlflow.pyfunc.loaders.chat_model
+4 | import mlflow.pyfunc.loaders.code_model
+  |
+help: Remove unused import: `mlflow.pyfunc.loaders.chat_agent`
+1 | 
+  - import mlflow.pyfunc.loaders.chat_agent
+2 | import mlflow.pyfunc.loaders.chat_model
+3 | import mlflow.pyfunc.loaders.code_model
+4 | from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+F401 [*] `mlflow.pyfunc.loaders.chat_model` imported but unused
+ --> f401_preview_submodule.py:3:8
+  |
+2 | import mlflow.pyfunc.loaders.chat_agent
+3 | import mlflow.pyfunc.loaders.chat_model
+  |        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+4 | import mlflow.pyfunc.loaders.code_model
+5 | from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+  |
+help: Remove unused import: `mlflow.pyfunc.loaders.chat_model`
+1 | 
+2 | import mlflow.pyfunc.loaders.chat_agent
+  - import mlflow.pyfunc.loaders.chat_model
+3 | import mlflow.pyfunc.loaders.code_model
+4 | from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+5 | 
+
+F401 [*] `mlflow.pyfunc.loaders.code_model` imported but unused
+ --> f401_preview_submodule.py:4:8
+  |
+2 | import mlflow.pyfunc.loaders.chat_agent
+3 | import mlflow.pyfunc.loaders.chat_model
+4 | import mlflow.pyfunc.loaders.code_model
+  |        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+5 | from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+  |
+help: Remove unused import: `mlflow.pyfunc.loaders.code_model`
+1 | 
+2 | import mlflow.pyfunc.loaders.chat_agent
+3 | import mlflow.pyfunc.loaders.chat_model
+  - import mlflow.pyfunc.loaders.code_model
+4 | from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+5 | 
+6 | if IS_PYDANTIC_V2_OR_NEWER:
+
+F401 [*] `mlflow.pyfunc.loaders.responses_agent` imported but unused
+ --> f401_preview_submodule.py:8:12
+  |
+7 | if IS_PYDANTIC_V2_OR_NEWER:
+8 |     import mlflow.pyfunc.loaders.responses_agent
+  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |
+help: Remove unused import: `mlflow.pyfunc.loaders.responses_agent`
+5 | from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+6 | 
+7 | if IS_PYDANTIC_V2_OR_NEWER:
+  -     import mlflow.pyfunc.loaders.responses_agent
+8 +     pass

crates/ruff_python_semantic/src/model.rs

@@ -1684,6 +1684,48 @@ impl<'a> SemanticModel<'a> {
         left == right
     }
 
+    /// Returns `true` if any execution path to `node` passes through `dominator`.
+    ///
+    /// More precisely, it returns true if the path of branches leading
+    /// to `dominator` is a prefix of the path of branches leading to `node`.
+    ///
+    /// In this code snippet:
+    ///
+    /// ```python
+    /// if cond:
+    ///     dominator
+    ///     if other_cond:
+    ///         node
+    /// else:
+    ///     other_node
+    /// ```
+    ///
+    /// we have that `node` is dominated by `dominator` but that
+    /// `other_node` is not dominated by `dominator`.
+    ///
+    /// This implementation assumes that the statements are in the same scope.
+    pub fn dominates(&self, dominator: NodeId, node: NodeId) -> bool {
+        // Collect the branch path for the left statement.
+        let dominator = self
+            .nodes
+            .branch_id(dominator)
+            .iter()
+            .flat_map(|branch_id| self.branches.ancestor_ids(*branch_id))
+            .collect::<Vec<_>>();
+
+        // Collect the branch path for the right statement.
+        let node = self
+            .nodes
+            .branch_id(node)
+            .iter()
+            .flat_map(|branch_id| self.branches.ancestor_ids(*branch_id))
+            .collect::<Vec<_>>();
+
+        // Note that the paths are in "reverse" order -
+        // from most nested to least nested.
+        node.ends_with(&dominator)
+    }
+
     /// Returns `true` if the given expression is an unused variable, or consists solely of
     /// references to other unused variables. This method is conservative in that it considers a
     /// variable to be "used" if it's shadowed by another variable with usages.