[pyflakes] Fix preview-mode bugs in F401 when attempting to autofix unused first-party submodule imports in an __init__.py file (#12569)

crates/ruff_linter/src/fix/edits.rs

@@ -1,7 +1,5 @@
 //! Interface for generating fix edits from higher-level actions (e.g., "remove an argument").
 
-use std::borrow::Cow;
-
 use anyhow::{Context, Result};
 
 use ruff_diagnostics::Edit;
@@ -126,7 +124,7 @@ pub(crate) fn remove_unused_imports<'a>(
 
 /// Edits to make the specified imports explicit, e.g. change `import x` to `import x as x`.
 pub(crate) fn make_redundant_alias<'a>(
-    member_names: impl Iterator<Item = Cow<'a, str>>,
+    member_names: impl Iterator<Item = &'a str>,
     stmt: &Stmt,
 ) -> Vec<Edit> {
     let aliases = match stmt {
@@ -527,7 +525,6 @@ fn all_lines_fit(
 #[cfg(test)]
 mod tests {
     use anyhow::{anyhow, Result};
-    use std::borrow::Cow;
     use test_case::test_case;
 
     use ruff_diagnostics::{Diagnostic, Edit, Fix};
@@ -619,7 +616,7 @@ x = 1 \
         let contents = "import x, y as y, z as bees";
         let stmt = parse_first_stmt(contents)?;
         assert_eq!(
-            make_redundant_alias(["x"].into_iter().map(Cow::from), &stmt),
+            make_redundant_alias(["x"].into_iter(), &stmt),
             vec![Edit::range_replacement(
                 String::from("x as x"),
                 TextRange::new(TextSize::new(7), TextSize::new(8)),
@@ -627,7 +624,7 @@ x = 1 \
             "make just one item redundant"
         );
         assert_eq!(
-            make_redundant_alias(vec!["x", "y"].into_iter().map(Cow::from), &stmt),
+            make_redundant_alias(vec!["x", "y"].into_iter(), &stmt),
             vec![Edit::range_replacement(
                 String::from("x as x"),
                 TextRange::new(TextSize::new(7), TextSize::new(8)),
@@ -635,7 +632,7 @@ x = 1 \
             "the second item is already a redundant alias"
         );
         assert_eq!(
-            make_redundant_alias(vec!["x", "z"].into_iter().map(Cow::from), &stmt),
+            make_redundant_alias(vec!["x", "z"].into_iter(), &stmt),
             vec![Edit::range_replacement(
                 String::from("x as x"),
                 TextRange::new(TextSize::new(7), TextSize::new(8)),

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

@@ -11,6 +11,7 @@ mod tests {
 
     use anyhow::Result;
     use regex::Regex;
+    use rustc_hash::FxHashMap;
 
     use test_case::test_case;
 
@@ -24,11 +25,12 @@ mod tests {
 
     use crate::linter::check_path;
     use crate::registry::{AsRule, Linter, Rule};
+    use crate::rules::isort;
     use crate::rules::pyflakes;
     use crate::settings::types::PreviewMode;
     use crate::settings::{flags, LinterSettings};
     use crate::source_kind::SourceKind;
-    use crate::test::{test_path, test_snippet};
+    use crate::test::{test_contents, test_path, test_snippet};
     use crate::{assert_messages, directives};
 
     #[test_case(Rule::UnusedImport, Path::new("F401_0.py"))]
@@ -232,6 +234,44 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(
+        r"import submodule.a",
+        "f401_preview_first_party_submodule_no_dunder_all"
+    )]
+    #[test_case(
+        r"
+        import submodule.a
+        __all__ = ['FOO']
+        FOO = 42",
+        "f401_preview_first_party_submodule_dunder_all"
+    )]
+    fn f401_preview_first_party_submodule(contents: &str, snapshot: &str) {
+        let diagnostics = test_contents(
+            &SourceKind::Python(dedent(contents).to_string()),
+            Path::new("f401_preview_first_party_submodule/__init__.py"),
+            &LinterSettings {
+                preview: PreviewMode::Enabled,
+                isort: isort::settings::Settings {
+                    // This case specifically tests the scenario where
+                    // the unused import is a first-party submodule import;
+                    // use the isort settings to ensure that the `submodule.a` import
+                    // is recognised as first-party in the test:
+                    known_modules: isort::categorize::KnownModules::new(
+                        vec!["submodule".parse().unwrap()],
+                        vec![],
+                        vec![],
+                        vec![],
+                        FxHashMap::default(),
+                    ),
+                    ..isort::settings::Settings::default()
+                },
+                ..LinterSettings::for_rule(Rule::UnusedImport)
+            },
+        )
+        .0;
+        assert_messages!(snapshot, diagnostics);
+    }
+
     #[test_case(Rule::UnusedImport, Path::new("F401_24/__init__.py"))]
     #[test_case(Rule::UnusedImport, Path::new("F401_25__all_nonempty/__init__.py"))]
     #[test_case(Rule::UnusedImport, Path::new("F401_26__all_empty/__init__.py"))]

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

@@ -9,7 +9,7 @@ use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast as ast;
 use ruff_python_ast::{Stmt, StmtImportFrom};
 use ruff_python_semantic::{
-    AnyImport, BindingKind, Exceptions, Imported, NodeId, Scope, SemanticModel,
+    AnyImport, BindingKind, Exceptions, Imported, NodeId, Scope, SemanticModel, SubmoduleImport,
 };
 use ruff_text_size::{Ranged, TextRange};
 
@@ -18,16 +18,6 @@ use crate::fix;
 use crate::registry::Rule;
 use crate::rules::{isort, isort::ImportSection, isort::ImportType};
 
-#[derive(Debug, Copy, Clone, Eq, PartialEq)]
-enum UnusedImportContext {
-    ExceptHandler,
-    Init {
-        first_party: bool,
-        dunder_all_count: usize,
-        ignore_init_module_imports: bool,
-    },
-}
-
 /// ## What it does
 /// Checks for unused imports.
 ///
@@ -111,8 +101,9 @@ pub struct UnusedImport {
     module: String,
     /// Name of the import binding
     binding: String,
-    context: Option<UnusedImportContext>,
+    context: UnusedImportContext,
     multiple: bool,
+    ignore_init_module_imports: bool,
 }
 
 impl Violation for UnusedImport {
@@ -122,17 +113,17 @@ impl Violation for UnusedImport {
     fn message(&self) -> String {
         let UnusedImport { name, context, .. } = self;
         match context {
-            Some(UnusedImportContext::ExceptHandler) => {
+            UnusedImportContext::ExceptHandler => {
                 format!(
                     "`{name}` imported but unused; consider using `importlib.util.find_spec` to test for availability"
                 )
             }
-            Some(UnusedImportContext::Init { .. }) => {
+            UnusedImportContext::DunderInitFirstParty { .. } => {
                 format!(
                     "`{name}` imported but unused; consider removing, adding to `__all__`, or using a redundant alias"
                 )
             }
-            None => format!("`{name}` imported but unused"),
+            UnusedImportContext::Other => format!("`{name}` imported but unused"),
         }
     }
 
@@ -142,30 +133,91 @@ impl Violation for UnusedImport {
             module,
             binding,
             multiple,
-            ..
+            ignore_init_module_imports,
+            context,
         } = self;
-        match self.context {
+        if *ignore_init_module_imports {
-            Some(UnusedImportContext::Init {
+            match context {
-                first_party: true,
+                UnusedImportContext::DunderInitFirstParty {
-                dunder_all_count: 1,
+                    dunder_all_count: DunderAllCount::Zero,
-                ignore_init_module_imports: true,
+                    submodule_import: false,
-            }) => Some(format!("Add unused import `{binding}` to __all__")),
+                } => return Some(format!("Use an explicit re-export: `{module} as {module}`")),
+                UnusedImportContext::DunderInitFirstParty {
+                    dunder_all_count: DunderAllCount::Zero,
+                    submodule_import: true,
+                } => {
+                    return Some(format!(
+                    "Use an explicit re-export: `import {parent} as {parent}; import {binding}`",
+                    parent = binding
+                        .split('.')
+                        .next()
+                        .expect("Expected all submodule imports to contain a '.'")
+                    ))
+                }
+                UnusedImportContext::DunderInitFirstParty {
+                    dunder_all_count: DunderAllCount::One,
+                    submodule_import: false,
+                } => return Some(format!("Add unused import `{binding}` to __all__")),
+                UnusedImportContext::DunderInitFirstParty {
+                    dunder_all_count: DunderAllCount::One,
+                    submodule_import: true,
+                } => {
+                    return Some(format!(
+                        "Add `{}` to __all__",
+                        binding
+                            .split('.')
+                            .next()
+                            .expect("Expected all submodule imports to contain a '.'")
+                    ))
+                }
+                UnusedImportContext::DunderInitFirstParty {
+                    dunder_all_count: DunderAllCount::Many,
+                    submodule_import: _,
+                }
+                | UnusedImportContext::ExceptHandler
+                | UnusedImportContext::Other => {}
+            }
+        }
+        Some(if *multiple {
+            "Remove unused import".to_string()
+        } else {
+            format!("Remove unused import: `{name}`")
+        })
+    }
+}
 
-            Some(UnusedImportContext::Init {
+/// Enumeration providing three possible answers to the question:
-                first_party: true,
+/// "How many `__all__` definitions are there in this file?"
-                dunder_all_count: 0,
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-                ignore_init_module_imports: true,
+enum DunderAllCount {
-            }) => Some(format!("Use an explicit re-export: `{module} as {module}`")),
+    Zero,
+    One,
+    Many,
+}
 
-            _ => Some(if *multiple {
+impl From<usize> for DunderAllCount {
-                "Remove unused import".to_string()
+    fn from(value: usize) -> Self {
-            } else {
+        match value {
-                format!("Remove unused import: `{name}`")
+            0 => Self::Zero,
-            }),
+            1 => Self::One,
+            _ => Self::Many,
         }
     }
 }
 
+#[derive(Debug, Copy, Clone, Eq, PartialEq, is_macro::Is)]
+enum UnusedImportContext {
+    /// The unused import occurs inside an except handler
+    ExceptHandler,
+    /// The unused import is a first-party import in an `__init__.py` file
+    DunderInitFirstParty {
+        dunder_all_count: DunderAllCount,
+        submodule_import: bool,
+    },
+    /// The unused import is something else
+    Other,
+}
+
 fn is_first_party(qualified_name: &str, level: u32, checker: &Checker) -> bool {
     let category = isort::categorize(
         qualified_name,
@@ -304,31 +356,20 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut
             .into_iter()
             .map(|binding| {
                 let context = if in_except_handler {
-                    Some(UnusedImportContext::ExceptHandler)
+                    UnusedImportContext::ExceptHandler
-                } else if in_init {
+                } else if in_init
-                    Some(UnusedImportContext::Init {
+                    && is_first_party(&binding.import.qualified_name().to_string(), level, checker)
-                        first_party: is_first_party(
+                {
-                            &binding.import.qualified_name().to_string(),
+                    UnusedImportContext::DunderInitFirstParty {
-                            level,
+                        dunder_all_count: DunderAllCount::from(dunder_all_exprs.len()),
-                            checker,
+                        submodule_import: binding.import.is_submodule_import(),
-                        ),
+                    }
-                        dunder_all_count: dunder_all_exprs.len(),
-                        ignore_init_module_imports: !fix_init,
-                    })
                 } else {
-                    None
+                    UnusedImportContext::Other
                 };
                 (binding, context)
             })
-            .partition(|(_, context)| {
+            .partition(|(_, context)| context.is_dunder_init_first_party() && preview_mode);
-                matches!(
-                    context,
-                    Some(UnusedImportContext::Init {
-                        first_party: true,
-                        ..
-                    })
-                ) && preview_mode
-            });
 
         // generate fixes that are shared across bindings in the statement
         let (fix_remove, fix_reexport) =
@@ -344,7 +385,7 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut
                     fix_by_reexporting(
                         checker,
                         import_statement,
-                        to_reexport.iter().map(|(b, _)| b).collect::<Vec<_>>(),
+                        to_reexport.iter().map(|(b, _)| b),
                         &dunder_all_exprs,
                     )
                     .ok(),
@@ -364,6 +405,7 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut
                     binding: binding.name.to_string(),
                     context,
                     multiple,
+                    ignore_init_module_imports: !fix_init,
                 },
                 binding.range,
             );
@@ -387,8 +429,9 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope, diagnostics: &mut
                 name: binding.import.qualified_name().to_string(),
                 module: binding.import.member_name().to_string(),
                 binding: binding.name.to_string(),
-                context: None,
+                context: UnusedImportContext::Other,
                 multiple: false,
+                ignore_init_module_imports: !fix_init,
             },
             binding.range,
         );
@@ -412,6 +455,31 @@ struct ImportBinding<'a> {
     parent_range: Option<TextRange>,
 }
 
+impl<'a> ImportBinding<'a> {
+    /// The symbol that is stored in the outer scope as a result of this import.
+    ///
+    /// For example:
+    /// - `import foo` => `foo` symbol stored in outer scope
+    /// - `import foo as bar` => `bar` symbol stored in outer scope
+    /// - `from foo import bar` => `bar` symbol stored in outer scope
+    /// - `from foo import bar as baz` => `baz` symbol stored in outer scope
+    /// - `import foo.bar` => `foo` symbol stored in outer scope
+    fn symbol_stored_in_outer_scope(&self) -> &str {
+        match &self.import {
+            AnyImport::FromImport(_) => self.name,
+            AnyImport::Import(_) => self.name,
+            AnyImport::SubmoduleImport(SubmoduleImport { qualified_name }) => {
+                qualified_name.segments().first().unwrap_or_else(|| {
+                    panic!(
+                        "Expected an import binding to have a non-empty qualified name;
+                        got {qualified_name}"
+                    )
+                })
+            }
+        }
+    }
+}
+
 impl Ranged for ImportBinding<'_> {
     fn range(&self) -> TextRange {
         self.range
@@ -461,29 +529,31 @@ fn fix_by_removing_imports<'a>(
 
 /// Generate a [`Fix`] to make bindings in a statement explicit, either by adding them to `__all__`
 /// or changing them from `import a` to `import a as a`.
-fn fix_by_reexporting(
+fn fix_by_reexporting<'a>(
     checker: &Checker,
     node_id: NodeId,
-    mut imports: Vec<&ImportBinding>,
+    imports: impl IntoIterator<Item = &'a ImportBinding<'a>>,
     dunder_all_exprs: &[&ast::Expr],
 ) -> Result<Fix> {
     let statement = checker.semantic().statement(node_id);
-    if imports.is_empty() {
-        bail!("Expected import bindings");
-    }
 
-    imports.sort_by_key(|b| b.name);
+    let imports = {
+        let mut imports: Vec<&str> = imports
+            .into_iter()
+            .map(ImportBinding::symbol_stored_in_outer_scope)
+            .collect();
+        if imports.is_empty() {
+            bail!("Expected import bindings");
+        }
+        imports.sort_unstable();
+        imports
+    };
 
     let edits = match dunder_all_exprs {
-        [] => fix::edits::make_redundant_alias(
+        [] => fix::edits::make_redundant_alias(imports.into_iter(), statement),
-            imports.iter().map(|b| b.import.member_name()),
+        [dunder_all] => {
-            statement,
+            fix::edits::add_to_dunder_all(imports.into_iter(), dunder_all, checker.stylist())
-        ),
+        }
-        [dunder_all] => fix::edits::add_to_dunder_all(
-            imports.iter().map(|b| b.name),
-            dunder_all,
-            checker.stylist(),
-        ),
         _ => bail!("Cannot offer a fix when there are multiple __all__ definitions"),
     };
 

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F401_deprecated_option_F401_24____init__.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:19:8: F401 [*] `sys` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:19:8: F401 [*] `sys` imported but unused
    |
 19 | import sys  # F401: remove unused
    |        ^^^ F401

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F401_deprecated_option_F401_25__all_nonempty____init__.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:19:8: F401 [*] `sys` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:19:8: F401 [*] `sys` imported but unused
    |
 19 | import sys  # F401: remove unused
    |        ^^^ F401

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F401_stable_F401_24____init__.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:19:8: F401 `sys` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:19:8: F401 `sys` imported but unused
    |
 19 | import sys  # F401: remove unused
    |        ^^^ F401

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F401_stable_F401_25__all_nonempty____init__.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:19:8: F401 `sys` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:19:8: F401 `sys` imported but unused
    |
 19 | import sys  # F401: remove unused
    |        ^^^ F401

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

@@ -0,0 +1,18 @@
+---
+source: crates/ruff_linter/src/rules/pyflakes/mod.rs
+---
+__init__.py:2:8: F401 [*] `submodule.a` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+  |
+2 | import submodule.a
+  |        ^^^^^^^^^^^ F401
+3 | __all__ = ['FOO']
+4 | FOO = 42
+  |
+  = help: Add `submodule` to __all__
+
+ℹ Safe fix
+1 1 | 
+2 2 | import submodule.a
+3   |-__all__ = ['FOO']
+  3 |+__all__ = ['FOO', 'submodule']
+4 4 | FOO = 42

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

@@ -0,0 +1,9 @@
+---
+source: crates/ruff_linter/src/rules/pyflakes/mod.rs
+---
+__init__.py:1:8: F401 `submodule.a` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+  |
+1 | import submodule.a
+  |        ^^^^^^^^^^^ F401
+  |
+  = help: Use an explicit re-export: `import submodule as submodule; import submodule.a`

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

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:1:8: F401 `os` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:1:8: F401 `os` imported but unused
   |
 1 | import os
   |        ^^ F401

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__preview__F401_F401_24____init__.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:19:8: F401 [*] `sys` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:19:8: F401 [*] `sys` imported but unused
    |
 19 | import sys  # F401: remove unused
    |        ^^^ F401

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__preview__F401_F401_25__all_nonempty____init__.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:19:8: F401 [*] `sys` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:19:8: F401 [*] `sys` imported but unused
    |
 19 | import sys  # F401: remove unused
    |        ^^^ F401

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__preview__F401___init__.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-__init__.py:1:8: F401 [*] `os` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
+__init__.py:1:8: F401 [*] `os` imported but unused
   |
 1 | import os
   |        ^^ F401