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

Author: AlexWaygood

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,23 +616,23 @@ 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)),
             )],
             "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)),
             )],
             "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"))]