F401 - Recommend adding unused import bindings to __all__

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_25__all/__init__.py → crates/ruff_linter/resources/test/fixtures/pyflakes/F401_25__all_nonempty/__init__.py

@@ -1,4 +1,4 @@
-"""__init__.py with __all__
+"""__init__.py with nonempty __all__
 
 Unused stdlib and third party imports are unsafe removals
 
@@ -33,10 +33,10 @@
 from . import exported  # Ok: is exported in __all__
 
 
-# from . import unused # F401: add to __all__
+from . import unused  # F401: add to __all__
 
 
-# from . import renamed as bees # F401: add to __all__
+from . import renamed as bees  # F401: add to __all__
 
 
 __all__ = ["argparse", "exported"]

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_26__all_empty/__init__.py

@@ -0,0 +1,11 @@
+"""__init__.py with empty __all__
+"""
+
+
+from . import unused  # F401: add to __all__
+
+
+from . import renamed as bees  # F401: add to __all__
+
+
+__all__ = []

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_26__all_empty/renamed.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_26__all_empty/unused.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_27__all_mistyped/__init__.py

@@ -0,0 +1,11 @@
+"""__init__.py with mis-typed __all__
+"""
+
+
+from . import unused  # F401: recommend add to all w/o fix
+
+
+from . import renamed as bees  # F401: recommend add to all w/o fix
+
+
+__all__ = None

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_27__all_mistyped/renamed.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_27__all_mistyped/unused.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_28__all_multiple/__init__.py

@@ -0,0 +1,8 @@
+"""__init__.py with multiple imports added to all in one edit
+"""
+
+
+from . import unused, renamed as bees  # F401: add to __all__
+
+
+__all__ = [];

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_28__all_multiple/renamed.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_28__all_multiple/unused.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_29__all_conditional/__init__.py

@@ -0,0 +1,16 @@
+"""__init__.py with __all__ populated by conditional plus-eq
+
+multiple __all__ so cannot offer a fix to add to them
+"""
+
+import sys
+
+from . import unused, exported, renamed as bees
+
+if sys.version_info > (3, 9):
+    from . import also_exported
+
+__all__ = ["exported"]
+
+if sys.version_info >= (3, 9):
+    __all__ += ["also_exported"]

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_29__all_conditional/also_exported.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_29__all_conditional/exported.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_29__all_conditional/renamed.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/resources/test/fixtures/pyflakes/F401_29__all_conditional/unused.py

@@ -0,0 +1 @@
+# empty module imported by __init__.py for test fixture

crates/ruff_linter/src/fix/edits.rs

@@ -1,10 +1,12 @@
 //! 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;
 use ruff_python_ast::parenthesize::parenthesized_range;
-use ruff_python_ast::{self as ast, Arguments, ExceptHandler, Stmt};
+use ruff_python_ast::{self as ast, Arguments, ExceptHandler, Expr, ExprList, Stmt};
 use ruff_python_ast::{AnyNodeRef, ArgOrKeyword};
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
@@ -124,7 +126,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 = &'a str>,
+    member_names: impl Iterator<Item = Cow<'a, str>>,
     stmt: &Stmt,
 ) -> Vec<Edit> {
     let aliases = match stmt {
@@ -144,6 +146,53 @@ pub(crate) fn make_redundant_alias<'a>(
         .collect()
 }
 
+/// Fix to add the specified imports to the `__all__` export list.
+pub(crate) fn add_to_dunder_all<'a>(
+    names: impl Iterator<Item = &'a str>,
+    expr: &Expr,
+    stylist: &Stylist,
+) -> Vec<Edit> {
+    let (insertion_point, export_prefix_length) = match expr {
+        Expr::List(ExprList { elts, range, .. }) => (
+            elts.last()
+                .map_or(range.end() - "]".text_len(), Ranged::end),
+            elts.len(),
+        ),
+        Expr::Tuple(tup) if tup.parenthesized => (
+            tup.elts
+                .last()
+                .map_or(tup.end() - ")".text_len(), Ranged::end),
+            tup.elts.len(),
+        ),
+        Expr::Tuple(tup) if !tup.parenthesized => (
+            tup.elts
+                .last()
+                .expect("unparenthesized empty tuple is not possible")
+                .range()
+                .end(),
+            tup.elts.len(),
+        ),
+        _ => {
+            // we don't know how to insert into this expression
+            return vec![];
+        }
+    };
+    let quote = stylist.quote();
+    let mut edits: Vec<_> = names
+        .enumerate()
+        .map(|(offset, name)| match export_prefix_length + offset {
+            0 => Edit::insertion(format!("{quote}{name}{quote}"), insertion_point),
+            _ => Edit::insertion(format!(", {quote}{name}{quote}"), insertion_point),
+        })
+        .collect();
+    if let Expr::Tuple(tup) = expr {
+        if tup.parenthesized && export_prefix_length + edits.len() == 1 {
+            edits.push(Edit::insertion(",".to_string(), insertion_point));
+        }
+    }
+    edits
+}
+
 #[derive(Debug, Copy, Clone)]
 pub(crate) enum Parentheses {
     /// Remove parentheses, if the removed argument is the only argument left.
@@ -477,14 +526,20 @@ fn all_lines_fit(
 
 #[cfg(test)]
 mod tests {
-    use anyhow::Result;
+    use anyhow::{anyhow, Result};
+    use std::borrow::Cow;
+    use test_case::test_case;
 
-    use ruff_diagnostics::Edit;
-    use ruff_python_parser::parse_suite;
+    use ruff_diagnostics::{Diagnostic, Edit, Fix};
+    use ruff_python_codegen::Stylist;
+    use ruff_python_parser::{lexer, parse_expression, parse_suite, Mode};
     use ruff_source_file::Locator;
     use ruff_text_size::{Ranged, TextRange, TextSize};
 
-    use crate::fix::edits::{make_redundant_alias, next_stmt_break, trailing_semicolon};
+    use crate::fix::apply_fixes;
+    use crate::fix::edits::{
+        add_to_dunder_all, make_redundant_alias, next_stmt_break, trailing_semicolon,
+    };
 
     #[test]
     fn find_semicolon() -> Result<()> {
@@ -562,28 +617,71 @@ x = 1 \
         let program = parse_suite(contents).unwrap();
         let stmt = program.first().unwrap();
         assert_eq!(
-            make_redundant_alias(["x"].into_iter(), stmt),
+            make_redundant_alias(["x"].into_iter().map(Cow::from), 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(), stmt),
+            make_redundant_alias(vec!["x", "y"].into_iter().map(Cow::from), 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(), stmt),
+            make_redundant_alias(vec!["x", "z"].into_iter().map(Cow::from), stmt),
             vec![Edit::range_replacement(
                 String::from("x as x"),
                 TextRange::new(TextSize::new(7), TextSize::new(8)),
             )],
             "the third item is already aliased to something else"
         );
     }
+
+    #[test_case("()",             &["x", "y"], r#"("x", "y")"#             ; "2 into empty tuple")]
+    #[test_case("()",             &["x"],      r#"("x",)"#                 ; "1 into empty tuple adding a trailing comma")]
+    #[test_case("[]",             &["x", "y"], r#"["x", "y"]"#             ; "2 into empty list")]
+    #[test_case("[]",             &["x"],      r#"["x"]"#                  ; "1 into empty list")]
+    #[test_case(r#""a", "b""#,    &["x", "y"], r#""a", "b", "x", "y""#     ; "2 into unparenthesized tuple")]
+    #[test_case(r#""a", "b""#,    &["x"],      r#""a", "b", "x""#          ; "1 into unparenthesized tuple")]
+    #[test_case(r#""a", "b","#,   &["x", "y"], r#""a", "b", "x", "y","#    ; "2 into unparenthesized tuple w/trailing comma")]
+    #[test_case(r#""a", "b","#,   &["x"],      r#""a", "b", "x","#         ; "1 into unparenthesized tuple w/trailing comma")]
+    #[test_case(r#"("a", "b")"#,  &["x", "y"], r#"("a", "b", "x", "y")"#   ; "2 into nonempty tuple")]
+    #[test_case(r#"("a", "b")"#,  &["x"],      r#"("a", "b", "x")"#        ; "1 into nonempty tuple")]
+    #[test_case(r#"("a", "b",)"#, &["x", "y"], r#"("a", "b", "x", "y",)"#  ; "2 into nonempty tuple w/trailing comma")]
+    #[test_case(r#"("a", "b",)"#, &["x"],      r#"("a", "b", "x",)"#       ; "1 into nonempty tuple w/trailing comma")]
+    #[test_case(r#"["a", "b",]"#, &["x", "y"], r#"["a", "b", "x", "y",]"#  ; "2 into nonempty list w/trailing comma")]
+    #[test_case(r#"["a", "b",]"#, &["x"],      r#"["a", "b", "x",]"#       ; "1 into nonempty list w/trailing comma")]
+    #[test_case(r#"["a", "b"]"#,  &["x", "y"], r#"["a", "b", "x", "y"]"#   ; "2 into nonempty list")]
+    #[test_case(r#"["a", "b"]"#,  &["x"],      r#"["a", "b", "x"]"#        ; "1 into nonempty list")]
+    fn add_to_dunder_all_test(raw: &str, names: &[&str], expect: &str) -> Result<()> {
+        let locator = Locator::new(raw);
+        let edits = {
+            let expr = parse_expression(raw)?;
+            let stylist = Stylist::from_tokens(
+                &lexer::lex(raw, Mode::Expression).collect::<Vec<_>>(),
+                &locator,
+            );
+            // SUT
+            add_to_dunder_all(names.iter().copied(), &expr, &stylist)
+        };
+        let diag = {
+            use crate::rules::pycodestyle::rules::MissingNewlineAtEndOfFile;
+            let mut iter = edits.into_iter();
+            Diagnostic::new(
+                MissingNewlineAtEndOfFile, // The choice of rule here is arbitrary.
+                TextRange::default(),
+            )
+            .with_fix(Fix::safe_edits(
+                iter.next().ok_or(anyhow!("expected edits nonempty"))?,
+                iter,
+            ))
+        };
+        assert_eq!(apply_fixes([diag].iter(), &locator).code, expect);
+        Ok(())
+    }
 }

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

@@ -207,7 +207,11 @@ mod tests {
     #[test_case(Rule::UnusedVariable, Path::new("F841_4.py"))]
     #[test_case(Rule::UnusedImport, Path::new("__init__.py"))]
     #[test_case(Rule::UnusedImport, Path::new("F401_24/__init__.py"))]
-    #[test_case(Rule::UnusedImport, Path::new("F401_25__all/__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"))]
+    #[test_case(Rule::UnusedImport, Path::new("F401_27__all_mistyped/__init__.py"))]
+    #[test_case(Rule::UnusedImport, Path::new("F401_28__all_multiple/__init__.py"))]
+    #[test_case(Rule::UnusedImport, Path::new("F401_29__all_conditional/__init__.py"))]
     fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!(
             "preview__{}_{}",