Skip to content

Commit

Permalink
Support length-2 lists in dictionary comprehension rewrites
Browse files Browse the repository at this point in the history
  • Loading branch information
@charliermarsh
charliermarsh committed Sep 3, 2023
1 parent 577280c commit c1c7e21
Show file tree
Hide file tree
Showing 3 changed files with 148 additions and 125 deletions.
2 changes: 2 additions & 0 deletions crates/ruff/resources/test/fixtures/flake8_comprehensions/C417.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
list(map(lambda x: x * 2, nums))
set(map(lambda x: x % 2 == 0, nums))
dict(map(lambda v: (v, v**2), nums))
dict(map(lambda v: [v, v**2], nums))
map(lambda: "const", nums)
map(lambda _: 3.0, nums)
_ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
Expand Down Expand Up @@ -40,3 +41,4 @@ def func(arg1: int, arg2: int = 4):

# Ok because multiple arguments are allowed.
dict(map(lambda k, v: (k, v), keys, values))

28 changes: 14 additions & 14 deletions crates/ruff/src/rules/flake8_comprehensions/fixes.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -966,21 +966,21 @@ pub(crate) fn fix_unnecessary_map(
}));
}
ObjectType::Dict => {
let (key, value) = if let Expression::Tuple(tuple) = func_body.body.as_ref() {
if tuple.elements.len() != 2 {
bail!("Expected two elements")
let elements = match func_body.body.as_ref() {
Expression::Tuple(tuple) => &tuple.elements,
Expression::List(list) => &list.elements,
_ => {
bail!("Expected tuple or list for dictionary comprehension")
}

let Some(Element::Simple { value: key, .. }) = &tuple.elements.get(0) else {
bail!("Expected tuple to contain a key as the first element");
};
let Some(Element::Simple { value, .. }) = &tuple.elements.get(1) else {
bail!("Expected tuple to contain a key as the second element");
};

(key, value)
} else {
bail!("Expected tuple for dict comprehension")
};
let [key, value] = elements.as_slice() else {
bail!("Expected container to include two elements");
};
let Element::Simple { value: key, .. } = key else {
bail!("Expected container to use a key as the first element");
};
let Element::Simple { value, .. } = value else {
bail!("Expected container to use a value as the second element");
};

tree = Expression::DictComp(Box::new(DictComp {
Expand Down
243 changes: 132 additions & 111 deletions ...ke8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C417_C417.py.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ C417.py:5:1: C417 [*] Unnecessary `map` usage (rewrite using a `list` comprehens
5 |+[x * 2 for x in nums]
6 6 | set(map(lambda x: x % 2 == 0, nums))
7 7 | dict(map(lambda v: (v, v**2), nums))
8 8 | map(lambda: "const", nums)
8 8 | dict(map(lambda v: [v, v**2], nums))

C417.py:6:1: C417 [*] Unnecessary `map` usage (rewrite using a `set` comprehension)
|
Expand All @@ -70,7 +70,7 @@ C417.py:6:1: C417 [*] Unnecessary `map` usage (rewrite using a `set` comprehensi
6 | set(map(lambda x: x % 2 == 0, nums))
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
7 | dict(map(lambda v: (v, v**2), nums))
8 | map(lambda: "const", nums)
8 | dict(map(lambda v: [v, v**2], nums))
|
= help: Replace `map` with a `set` comprehension

Expand All @@ -81,17 +81,17 @@ C417.py:6:1: C417 [*] Unnecessary `map` usage (rewrite using a `set` comprehensi
6 |-set(map(lambda x: x % 2 == 0, nums))
6 |+{x % 2 == 0 for x in nums}
7 7 | dict(map(lambda v: (v, v**2), nums))
8 8 | map(lambda: "const", nums)
9 9 | map(lambda _: 3.0, nums)
8 8 | dict(map(lambda v: [v, v**2], nums))
9 9 | map(lambda: "const", nums)

C417.py:7:1: C417 [*] Unnecessary `map` usage (rewrite using a `dict` comprehension)
|
5 | list(map(lambda x: x * 2, nums))
6 | set(map(lambda x: x % 2 == 0, nums))
7 | dict(map(lambda v: (v, v**2), nums))
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
8 | map(lambda: "const", nums)
9 | map(lambda _: 3.0, nums)
8 | dict(map(lambda v: [v, v**2], nums))
9 | map(lambda: "const", nums)
|
= help: Replace `map` with a `dict` comprehension

Expand All @@ -101,172 +101,193 @@ C417.py:7:1: C417 [*] Unnecessary `map` usage (rewrite using a `dict` comprehens
6 6 | set(map(lambda x: x % 2 == 0, nums))
7 |-dict(map(lambda v: (v, v**2), nums))
7 |+{v: v**2 for v in nums}
8 8 | map(lambda: "const", nums)
9 9 | map(lambda _: 3.0, nums)
10 10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
8 8 | dict(map(lambda v: [v, v**2], nums))
9 9 | map(lambda: "const", nums)
10 10 | map(lambda _: 3.0, nums)

C417.py:8:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
C417.py:8:1: C417 [*] Unnecessary `map` usage (rewrite using a `dict` comprehension)
|
6 | set(map(lambda x: x % 2 == 0, nums))
7 | dict(map(lambda v: (v, v**2), nums))
8 | map(lambda: "const", nums)
| ^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
9 | map(lambda _: 3.0, nums)
10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
8 | dict(map(lambda v: [v, v**2], nums))
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
9 | map(lambda: "const", nums)
10 | map(lambda _: 3.0, nums)
|
= help: Replace `map` with a generator expression
= help: Replace `map` with a `dict` comprehension

ℹ Suggested fix
5 5 | list(map(lambda x: x * 2, nums))
6 6 | set(map(lambda x: x % 2 == 0, nums))
7 7 | dict(map(lambda v: (v, v**2), nums))
8 |-map(lambda: "const", nums)
8 |+("const" for _ in nums)
9 9 | map(lambda _: 3.0, nums)
10 10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 11 | all(map(lambda v: isinstance(v, dict), nums))
8 |-dict(map(lambda v: [v, v**2], nums))
8 |+{v: v**2 for v in nums}
9 9 | map(lambda: "const", nums)
10 10 | map(lambda _: 3.0, nums)
11 11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))

C417.py:9:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
|
7 | dict(map(lambda v: (v, v**2), nums))
8 | map(lambda: "const", nums)
9 | map(lambda _: 3.0, nums)
| ^^^^^^^^^^^^^^^^^^^^^^^^ C417
10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 | all(map(lambda v: isinstance(v, dict), nums))
8 | dict(map(lambda v: [v, v**2], nums))
9 | map(lambda: "const", nums)
| ^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
10 | map(lambda _: 3.0, nums)
11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
|
= help: Replace `map` with a generator expression

ℹ Suggested fix
6 6 | set(map(lambda x: x % 2 == 0, nums))
7 7 | dict(map(lambda v: (v, v**2), nums))
8 8 | map(lambda: "const", nums)
9 |-map(lambda _: 3.0, nums)
9 |+(3.0 for _ in nums)
10 10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 11 | all(map(lambda v: isinstance(v, dict), nums))
12 12 | filter(func, map(lambda v: v, nums))
8 8 | dict(map(lambda v: [v, v**2], nums))
9 |-map(lambda: "const", nums)
9 |+("const" for _ in nums)
10 10 | map(lambda _: 3.0, nums)
11 11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
12 12 | all(map(lambda v: isinstance(v, dict), nums))

C417.py:10:13: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
C417.py:10:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
|
8 | map(lambda: "const", nums)
9 | map(lambda _: 3.0, nums)
10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
11 | all(map(lambda v: isinstance(v, dict), nums))
12 | filter(func, map(lambda v: v, nums))
8 | dict(map(lambda v: [v, v**2], nums))
9 | map(lambda: "const", nums)
10 | map(lambda _: 3.0, nums)
| ^^^^^^^^^^^^^^^^^^^^^^^^ C417
11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
12 | all(map(lambda v: isinstance(v, dict), nums))
|
= help: Replace `map` with a generator expression

ℹ Suggested fix
7 7 | dict(map(lambda v: (v, v**2), nums))
8 8 | map(lambda: "const", nums)
9 9 | map(lambda _: 3.0, nums)
10 |-_ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
10 |+_ = "".join((x in nums and "1" or "0" for x in range(123)))
11 11 | all(map(lambda v: isinstance(v, dict), nums))
12 12 | filter(func, map(lambda v: v, nums))
13 13 |
8 8 | dict(map(lambda v: [v, v**2], nums))
9 9 | map(lambda: "const", nums)
10 |-map(lambda _: 3.0, nums)
10 |+(3.0 for _ in nums)
11 11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
12 12 | all(map(lambda v: isinstance(v, dict), nums))
13 13 | filter(func, map(lambda v: v, nums))

C417.py:11:13: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
|
9 | map(lambda: "const", nums)
10 | map(lambda _: 3.0, nums)
11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
12 | all(map(lambda v: isinstance(v, dict), nums))
13 | filter(func, map(lambda v: v, nums))
|
= help: Replace `map` with a generator expression

ℹ Suggested fix
8 8 | dict(map(lambda v: [v, v**2], nums))
9 9 | map(lambda: "const", nums)
10 10 | map(lambda _: 3.0, nums)
11 |-_ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 |+_ = "".join((x in nums and "1" or "0" for x in range(123)))
12 12 | all(map(lambda v: isinstance(v, dict), nums))
13 13 | filter(func, map(lambda v: v, nums))
14 14 |

C417.py:11:5: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
C417.py:12:5: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
|
9 | map(lambda _: 3.0, nums)
10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 | all(map(lambda v: isinstance(v, dict), nums))
10 | map(lambda _: 3.0, nums)
11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
12 | all(map(lambda v: isinstance(v, dict), nums))
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
12 | filter(func, map(lambda v: v, nums))
13 | filter(func, map(lambda v: v, nums))
|
= help: Replace `map` with a generator expression

ℹ Suggested fix
8 8 | map(lambda: "const", nums)
9 9 | map(lambda _: 3.0, nums)
10 10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 |-all(map(lambda v: isinstance(v, dict), nums))
11 |+all((isinstance(v, dict) for v in nums))
12 12 | filter(func, map(lambda v: v, nums))
13 13 |
14 14 | # When inside f-string, then the fix should be surrounded by whitespace
9 9 | map(lambda: "const", nums)
10 10 | map(lambda _: 3.0, nums)
11 11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
12 |-all(map(lambda v: isinstance(v, dict), nums))
12 |+all((isinstance(v, dict) for v in nums))
13 13 | filter(func, map(lambda v: v, nums))
14 14 |
15 15 | # When inside f-string, then the fix should be surrounded by whitespace

C417.py:12:14: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
C417.py:13:14: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
|
10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 | all(map(lambda v: isinstance(v, dict), nums))
12 | filter(func, map(lambda v: v, nums))
11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
12 | all(map(lambda v: isinstance(v, dict), nums))
13 | filter(func, map(lambda v: v, nums))
| ^^^^^^^^^^^^^^^^^^^^^^ C417
13 |
14 | # When inside f-string, then the fix should be surrounded by whitespace
14 |
15 | # When inside f-string, then the fix should be surrounded by whitespace
|
= help: Replace `map` with a generator expression

ℹ Suggested fix
9 9 | map(lambda _: 3.0, nums)
10 10 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
11 11 | all(map(lambda v: isinstance(v, dict), nums))
12 |-filter(func, map(lambda v: v, nums))
12 |+filter(func, (v for v in nums))
13 13 |
14 14 | # When inside f-string, then the fix should be surrounded by whitespace
15 15 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"
10 10 | map(lambda _: 3.0, nums)
11 11 | _ = "".join(map(lambda x: x in nums and "1" or "0", range(123)))
12 12 | all(map(lambda v: isinstance(v, dict), nums))
13 |-filter(func, map(lambda v: v, nums))
13 |+filter(func, (v for v in nums))
14 14 |
15 15 | # When inside f-string, then the fix should be surrounded by whitespace
16 16 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"

C417.py:15:8: C417 [*] Unnecessary `map` usage (rewrite using a `set` comprehension)
C417.py:16:8: C417 [*] Unnecessary `map` usage (rewrite using a `set` comprehension)
|
14 | # When inside f-string, then the fix should be surrounded by whitespace
15 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"
15 | # When inside f-string, then the fix should be surrounded by whitespace
16 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
16 | _ = f"{dict(map(lambda v: (v, v**2), nums))}"
17 | _ = f"{dict(map(lambda v: (v, v**2), nums))}"
|
= help: Replace `map` with a `set` comprehension

ℹ Suggested fix
12 12 | filter(func, map(lambda v: v, nums))
13 13 |
14 14 | # When inside f-string, then the fix should be surrounded by whitespace
15 |-_ = f"{set(map(lambda x: x % 2 == 0, nums))}"
15 |+_ = f"{ {x % 2 == 0 for x in nums} }"
16 16 | _ = f"{dict(map(lambda v: (v, v**2), nums))}"
17 17 |
18 18 | # False negatives.
13 13 | filter(func, map(lambda v: v, nums))
14 14 |
15 15 | # When inside f-string, then the fix should be surrounded by whitespace
16 |-_ = f"{set(map(lambda x: x % 2 == 0, nums))}"
16 |+_ = f"{ {x % 2 == 0 for x in nums} }"
17 17 | _ = f"{dict(map(lambda v: (v, v**2), nums))}"
18 18 |
19 19 | # False negatives.

C417.py:16:8: C417 [*] Unnecessary `map` usage (rewrite using a `dict` comprehension)
C417.py:17:8: C417 [*] Unnecessary `map` usage (rewrite using a `dict` comprehension)
|
14 | # When inside f-string, then the fix should be surrounded by whitespace
15 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"
16 | _ = f"{dict(map(lambda v: (v, v**2), nums))}"
15 | # When inside f-string, then the fix should be surrounded by whitespace
16 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"
17 | _ = f"{dict(map(lambda v: (v, v**2), nums))}"
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
17 |
18 | # False negatives.
18 |
19 | # False negatives.
|
= help: Replace `map` with a `dict` comprehension

ℹ Suggested fix
13 13 |
14 14 | # When inside f-string, then the fix should be surrounded by whitespace
15 15 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"
16 |-_ = f"{dict(map(lambda v: (v, v**2), nums))}"
16 |+_ = f"{ {v: v**2 for v in nums} }"
17 17 |
18 18 | # False negatives.
19 19 | map(lambda x=2, y=1: x + y, nums, nums)
14 14 |
15 15 | # When inside f-string, then the fix should be surrounded by whitespace
16 16 | _ = f"{set(map(lambda x: x % 2 == 0, nums))}"
17 |-_ = f"{dict(map(lambda v: (v, v**2), nums))}"
17 |+_ = f"{ {v: v**2 for v in nums} }"
18 18 |
19 19 | # False negatives.
20 20 | map(lambda x=2, y=1: x + y, nums, nums)

C417.py:34:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
C417.py:35:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression)
|
33 | # Error: the `x` is overridden by the inner lambda.
34 | map(lambda x: lambda x: x, range(4))
34 | # Error: the `x` is overridden by the inner lambda.
35 | map(lambda x: lambda x: x, range(4))
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417
35 |
36 | # Ok because of the default parameters, and variadic arguments.
36 |
37 | # Ok because of the default parameters, and variadic arguments.
|
= help: Replace `map` with a generator expression

ℹ Suggested fix
31 31 | map(lambda x: lambda: x, range(4))
32 32 |
33 33 | # Error: the `x` is overridden by the inner lambda.
34 |-map(lambda x: lambda x: x, range(4))
34 |+(lambda x: x for x in range(4))
35 35 |
36 36 | # Ok because of the default parameters, and variadic arguments.
37 37 | map(lambda x=1: x, nums)
32 32 | map(lambda x: lambda: x, range(4))
33 33 |
34 34 | # Error: the `x` is overridden by the inner lambda.
35 |-map(lambda x: lambda x: x, range(4))
35 |+(lambda x: x for x in range(4))
36 36 |
37 37 | # Ok because of the default parameters, and variadic arguments.
38 38 | map(lambda x=1: x, nums)


0 comments on commit c1c7e21

Please sign in to comment.