From e0df62b841500bd60b27ada7e9169aeb48036bc2 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 16 Mar 2023 20:54:29 -0400 Subject: [PATCH] Rewrite mock import with starred imports (#3566) --- .../test/fixtures/pyupgrade/UP026.py | 31 ++- .../pyupgrade/rules/rewrite_mock_import.rs | 91 ++++--- ...ff__rules__pyupgrade__tests__UP026.py.snap | 226 ++++++++++-------- 3 files changed, 194 insertions(+), 154 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP026.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP026.py index badd19a210b24..b61b4c299ca8e 100644 --- a/crates/ruff/resources/test/fixtures/pyupgrade/UP026.py +++ b/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() diff --git a/crates/ruff/src/rules/pyupgrade/rules/rewrite_mock_import.rs b/crates/ruff/src/rules/pyupgrade/rules/rewrite_mock_import.rs index b1b91225916c9..924dcc4ae0d96 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/rewrite_mock_import.rs +++ b/crates/ruff/src/rules/pyupgrade/rules/rewrite_mock_import.rs @@ -93,22 +93,6 @@ fn clean_import_aliases(aliases: Vec) -> (Vec, Vec 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>, 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 diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP026.py.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP026.py.snap index 3a5641049c81d..49a288bdd1e0b 100644 --- a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP026.py.snap +++ b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP026.py.snap @@ -28,18 +28,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 6 + row: 7 column: 11 end_location: - row: 6 + row: 7 column: 15 fix: content: "import sys\n from unittest import mock" location: - row: 6 + row: 7 column: 4 end_location: - row: 6 + row: 7 column: 20 parent: ~ - kind: @@ -48,18 +48,38 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 9 + row: 11 + column: 4 + end_location: + row: 11 + column: 22 + fix: + content: from unittest.mock import * + location: + row: 11 + column: 4 + end_location: + row: 11 + column: 22 + parent: ~ +- kind: + name: RewriteMockImport + body: "`mock` is deprecated, use `unittest.mock`" + suggestion: "Import from `unittest.mock` instead" + fixable: true + location: + row: 14 column: 7 end_location: - row: 9 + row: 14 column: 16 fix: content: from unittest import mock location: - row: 9 + row: 14 column: 0 end_location: - row: 9 + row: 14 column: 16 parent: ~ - kind: @@ -68,18 +88,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 12 + row: 17 column: 19 end_location: - row: 12 + row: 17 column: 23 fix: content: "import contextlib, sys\nfrom unittest import mock" location: - row: 12 + row: 17 column: 0 end_location: - row: 12 + row: 17 column: 28 parent: ~ - kind: @@ -88,18 +108,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 15 + row: 20 column: 7 end_location: - row: 15 + row: 20 column: 11 fix: content: "import sys\nfrom unittest import mock" location: - row: 15 + row: 20 column: 0 end_location: - row: 15 + row: 20 column: 16 parent: ~ - kind: @@ -108,18 +128,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 19 + row: 24 column: 0 end_location: - row: 19 + row: 24 column: 21 fix: content: from unittest import mock location: - row: 19 + row: 24 column: 0 end_location: - row: 19 + row: 24 column: 21 parent: ~ - kind: @@ -128,18 +148,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 22 + row: 27 column: 0 end_location: - row: 27 + row: 32 column: 1 fix: - content: "from unittest.mock import (\n a,\n b,\n c,\n)" + content: "from unittest.mock import (\n a,\n b,\n c,\n)\nfrom unittest import mock" location: - row: 22 + row: 27 column: 0 end_location: - row: 27 + row: 32 column: 1 parent: ~ - kind: @@ -148,18 +168,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 28 + row: 33 column: 0 end_location: - row: 33 + row: 38 column: 1 fix: - content: "from unittest.mock import (\n a,\n b,\n c,\n)" + content: "from unittest.mock import (\n a,\n b,\n c,\n)\nfrom unittest import mock" location: - row: 28 + row: 33 column: 0 end_location: - row: 33 + row: 38 column: 1 parent: ~ - kind: @@ -168,18 +188,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 36 + row: 41 column: 0 end_location: - row: 41 + row: 46 column: 1 fix: - content: "from unittest.mock import (\n a,\n b,\n c\n)" + content: "from unittest.mock import (\n a,\n b,\n c\n)\nfrom unittest import mock" location: - row: 36 + row: 41 column: 0 end_location: - row: 41 + row: 46 column: 1 parent: ~ - kind: @@ -188,18 +208,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 42 + row: 47 column: 0 end_location: - row: 47 + row: 52 column: 1 fix: - content: "from unittest.mock import (\n a,\n b,\n c\n)" + content: "from unittest.mock import (\n a,\n b,\n c\n)\nfrom unittest import mock" location: - row: 42 + row: 47 column: 0 end_location: - row: 47 + row: 52 column: 1 parent: ~ - kind: @@ -208,18 +228,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 48 + row: 53 column: 0 end_location: - row: 48 + row: 53 column: 30 fix: - content: "from unittest.mock import a, b, c" + content: "from unittest.mock import a, b, c\nfrom unittest import mock" location: - row: 48 + row: 53 column: 0 end_location: - row: 48 + row: 53 column: 30 parent: ~ - kind: @@ -228,18 +248,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 49 + row: 54 column: 0 end_location: - row: 49 + row: 54 column: 30 fix: - content: "from unittest.mock import a, b, c" + content: "from unittest.mock import a, b, c\nfrom unittest import mock" location: - row: 49 + row: 54 column: 0 end_location: - row: 49 + row: 54 column: 30 parent: ~ - kind: @@ -248,18 +268,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 53 + row: 58 column: 8 end_location: - row: 58 + row: 63 column: 9 fix: - content: "from unittest.mock import (\n a,\n b,\n c\n )" + content: "from unittest.mock import (\n a,\n b,\n c\n )\n from unittest import mock" location: - row: 53 + row: 58 column: 8 end_location: - row: 58 + row: 63 column: 9 parent: ~ - kind: @@ -268,18 +288,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 64 + row: 69 column: 7 end_location: - row: 64 + row: 69 column: 11 fix: content: "from unittest import mock\nfrom unittest import mock" location: - row: 64 + row: 69 column: 0 end_location: - row: 64 + row: 69 column: 17 parent: ~ - kind: @@ -288,18 +308,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 64 + row: 69 column: 13 end_location: - row: 64 + row: 69 column: 17 fix: content: "from unittest import mock\nfrom unittest import mock" location: - row: 64 + row: 69 column: 0 end_location: - row: 64 + row: 69 column: 17 parent: ~ - kind: @@ -308,18 +328,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 67 + row: 72 column: 7 end_location: - row: 67 + row: 72 column: 18 fix: content: from unittest import mock as foo location: - row: 67 + row: 72 column: 0 end_location: - row: 67 + row: 72 column: 18 parent: ~ - kind: @@ -328,18 +348,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 70 + row: 75 column: 0 end_location: - row: 70 + row: 75 column: 28 fix: content: from unittest import mock as foo location: - row: 70 + row: 75 column: 0 end_location: - row: 70 + row: 75 column: 28 parent: ~ - kind: @@ -348,18 +368,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 74 + row: 79 column: 11 end_location: - row: 74 + row: 79 column: 22 fix: content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock" location: - row: 74 + row: 79 column: 4 end_location: - row: 74 + row: 79 column: 41 parent: ~ - kind: @@ -368,18 +388,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 74 + row: 79 column: 24 end_location: - row: 74 + row: 79 column: 35 fix: content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock" location: - row: 74 + row: 79 column: 4 end_location: - row: 74 + row: 79 column: 41 parent: ~ - kind: @@ -388,18 +408,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 74 + row: 79 column: 37 end_location: - row: 74 + row: 79 column: 41 fix: content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock" location: - row: 74 + row: 79 column: 4 end_location: - row: 74 + row: 79 column: 41 parent: ~ - kind: @@ -408,18 +428,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 77 + row: 82 column: 11 end_location: - row: 77 + row: 82 column: 22 fix: content: "import os\n from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock" location: - row: 77 + row: 82 column: 4 end_location: - row: 77 + row: 82 column: 45 parent: ~ - kind: @@ -428,18 +448,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 77 + row: 82 column: 24 end_location: - row: 77 + row: 82 column: 35 fix: content: "import os\n from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock" location: - row: 77 + row: 82 column: 4 end_location: - row: 77 + row: 82 column: 45 parent: ~ - kind: @@ -448,18 +468,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 77 + row: 82 column: 37 end_location: - row: 77 + row: 82 column: 41 fix: content: "import os\n from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock" location: - row: 77 + row: 82 column: 4 end_location: - row: 77 + row: 82 column: 45 parent: ~ - kind: @@ -468,18 +488,18 @@ expression: diagnostics suggestion: "Import from `unittest.mock` instead" fixable: true location: - row: 81 + row: 86 column: 4 end_location: - row: 81 + row: 86 column: 51 fix: content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock" location: - row: 81 + row: 86 column: 4 end_location: - row: 81 + row: 86 column: 51 parent: ~ - kind: @@ -488,18 +508,18 @@ expression: diagnostics suggestion: "Replace `mock.mock` with `mock`" fixable: true location: - row: 88 + row: 93 column: 4 end_location: - row: 88 + row: 93 column: 13 fix: content: mock location: - row: 88 + row: 93 column: 4 end_location: - row: 88 + row: 93 column: 13 parent: ~