Rewrite mock import with starred imports (#3566)

crates/ruff/resources/test/fixtures/pyupgrade/UP026.py

@@ -1,24 +1,29 @@
-# These should be changed
+# Error (`from unittest import mock`)
 if True:
     import mock
 
+# Error (`from unittest import mock`)
 if True:
     import mock, sys
 
-# This goes to from unitest import mock
+# Error (`from unittest.mock import *`)
+if True:
+    from mock import *
+
+# Error (`from unittest import mock`)
 import mock.mock
 
-# Mock should go on a new line as `from unittest import mock`
+# Error (`from unittest import mock`)
 import contextlib, mock, sys
 
-# Mock should go on a new line as `from unittest import mock`
+# Error (`from unittest import mock`)
 import mock, sys
 x = "This code should be preserved one line below the mock"
 
-# Mock should go on a new line as `from unittest import mock`
+# Error (`from unittest import mock`)
 from mock import mock
 
-# Should keep trailing comma
+# Error (keep trailing comma)
 from mock import (
     mock,
     a,
@@ -32,7 +37,7 @@
     mock,
 )
 
-# Should not get a trailing comma
+# Error (avoid trailing comma)
 from mock import (
     mock,
     a,
@@ -57,16 +62,16 @@
             c
         )
 
-# These should not change:
+# OK
 import os, io
 
-# Mock should go on a new line as `from unittest import mock`
+# Error (`from unittest import mock`)
 import mock, mock
 
-# Mock should go on a new line as `from unittest import mock as foo`
+# Error (`from unittest import mock as foo`)
 import mock as foo
 
-# Mock should go on a new line as `from unittest import mock as foo`
+# Error (`from unittest import mock as foo`)
 from mock import mock as foo
 
 if True:
@@ -81,8 +86,8 @@
     from mock import mock as foo, mock as bar, mock
 
 
-# This should be unchanged.
+# OK.
 x = mock.Mock()
 
-# This should change to `mock.Mock`().
+# Error (`mock.Mock()`).
 x = mock.mock.Mock()

crates/ruff/src/rules/pyupgrade/rules/rewrite_mock_import.rs

@@ -93,22 +93,6 @@ fn clean_import_aliases(aliases: Vec<ImportAlias>) -> (Vec<ImportAlias>, Vec<Opt
     (clean_aliases, mock_aliases)
 }
 
-/// Return `true` if the aliases contain `mock`.
-fn includes_mock_member(aliases: &[ImportAlias]) -> bool {
-    for alias in aliases {
-        let ImportAlias { name, .. } = &alias;
-        // Ex) `import mock.mock`
-        if let NameOrAttribute::A(attribute_struct) = name {
-            if let Expression::Name(name_struct) = &*attribute_struct.value {
-                if name_struct.value == "mock" && attribute_struct.attr.value == "mock" {
-                    return true;
-                }
-            }
-        }
-    }
-    false
-}
-
 fn format_mocks(aliases: Vec<Option<AsName>>, indent: &str, stylist: &Stylist) -> String {
     let mut content = String::new();
     for alias in aliases {
@@ -180,20 +164,12 @@ fn format_import_from(
     let mut tree = match_module(module_text).unwrap();
     let mut import = match_import_from(&mut tree)?;
 
-    let ImportFrom {
-        names: ImportNames::Aliases(aliases),
+    if let ImportFrom {
+        names: ImportNames::Star(..),
         ..
-    } = import.clone() else {
-        unreachable!("Expected ImportNames::Aliases");
-    };
-
-    let has_mock_member = includes_mock_member(&aliases);
-    let (clean_aliases, mock_aliases) = clean_import_aliases(aliases);
-
-    Ok(if clean_aliases.is_empty() {
-        format_mocks(mock_aliases, indent, stylist)
-    } else {
-        import.names = ImportNames::Aliases(clean_aliases);
+    } = import
+    {
+        // Ex) `from mock import *`
         import.module = Some(NameOrAttribute::A(Box::new(Attribute {
             value: Box::new(Expression::Name(Box::new(Name {
                 value: "unittest",
@@ -212,22 +188,61 @@ fn format_import_from(
             lpar: vec![],
             rpar: vec![],
         })));
-
         let mut state = CodegenState {
             default_newline: stylist.line_ending(),
             default_indent: stylist.indentation(),
             ..CodegenState::default()
         };
         tree.codegen(&mut state);
+        Ok(state.to_string())
+    } else if let ImportFrom {
+        names: ImportNames::Aliases(aliases),
+        ..
+    } = import
+    {
+        // Ex) `from mock import mock`
+        let (clean_aliases, mock_aliases) = clean_import_aliases(aliases.clone());
+        Ok(if clean_aliases.is_empty() {
+            format_mocks(mock_aliases, indent, stylist)
+        } else {
+            import.names = ImportNames::Aliases(clean_aliases);
+            import.module = Some(NameOrAttribute::A(Box::new(Attribute {
+                value: Box::new(Expression::Name(Box::new(Name {
+                    value: "unittest",
+                    lpar: vec![],
+                    rpar: vec![],
+                }))),
+                attr: Name {
+                    value: "mock",
+                    lpar: vec![],
+                    rpar: vec![],
+                },
+                dot: Dot {
+                    whitespace_before: ParenthesizableWhitespace::default(),
+                    whitespace_after: ParenthesizableWhitespace::default(),
+                },
+                lpar: vec![],
+                rpar: vec![],
+            })));
 
-        let mut content = state.to_string();
-        if has_mock_member {
-            content.push_str(stylist.line_ending());
-            content.push_str(indent);
-            content.push_str(&format_mocks(mock_aliases, indent, stylist));
-        }
-        content
-    })
+            let mut state = CodegenState {
+                default_newline: stylist.line_ending(),
+                default_indent: stylist.indentation(),
+                ..CodegenState::default()
+            };
+            tree.codegen(&mut state);
+
+            let mut content = state.to_string();
+            if !mock_aliases.is_empty() {
+                content.push_str(stylist.line_ending());
+                content.push_str(indent);
+                content.push_str(&format_mocks(mock_aliases, indent, stylist));
+            }
+            content
+        })
+    } else {
+        unreachable!("Expected ImportNames::Aliases | ImportNames::Star");
+    }
 }
 
 /// UP026