feat: update C416 with dict comprehension (autofixable) (#3605)

astral-sh · Mar 19, 2023 · 3a65af4 · 3a65af4
1 parent 474aa0b
commit 3a65af4
Show file tree

Hide file tree

Showing 12 changed files with 207 additions and 47 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_comprehensions/C416.py b/crates/ruff/resources/test/fixtures/flake8_comprehensions/C416.py
@@ -1,6 +1,19 @@
 x = [1, 2, 3]
+y = [("a", 1), ("b", 2), ("c", 3)]
+z = [(1,), (2,), (3,)]
+d = {"a": 1, "b": 2, "c": 3}
+
 [i for i in x]
 {i for i in x}
+{k: v for k, v in y}
+{k: v for k, v in d.items()}
 
+[i for i, in z]
+[i for i, j in y]
 [i for i in x if i > 1]
 [i for i in x for j in x]
+
+{v: k for k, v in y}
+{k.foo: k for k in y}
+{k["foo"]: k for k in y}
+{k: v if v else None for k, v in y}
diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
@@ -3326,7 +3326,7 @@ where
             }
             ExprKind::ListComp { elt, generators } | ExprKind::SetComp { elt, generators } => {
                 if self.settings.rules.enabled(Rule::UnnecessaryComprehension) {
-                    flake8_comprehensions::rules::unnecessary_comprehension(
+                    flake8_comprehensions::rules::unnecessary_list_set_comprehension(
                         self, expr, elt, generators,
                     );
                 }
@@ -3335,7 +3335,22 @@ where
                 }
                 self.ctx.push_scope(ScopeKind::Generator);
             }
-            ExprKind::GeneratorExp { .. } | ExprKind::DictComp { .. } => {
+            ExprKind::DictComp {
+                key,
+                value,
+                generators,
+            } => {
+                if self.settings.rules.enabled(Rule::UnnecessaryComprehension) {
+                    flake8_comprehensions::rules::unnecessary_dict_comprehension(
+                        self, expr, key, value, generators,
+                    );
+                }
+                if self.settings.rules.enabled(Rule::FunctionUsesLoopVariable) {
+                    flake8_bugbear::rules::function_uses_loop_variable(self, &Node::Expr(expr));
+                }
+                self.ctx.push_scope(ScopeKind::Generator);
+            }
+            ExprKind::GeneratorExp { .. } => {
                 if self.settings.rules.enabled(Rule::FunctionUsesLoopVariable) {
                     flake8_bugbear::rules::function_uses_loop_variable(self, &Node::Expr(expr));
                 }

diff --git a/crates/ruff/src/rules/flake8_comprehensions/fixes.rs b/crates/ruff/src/rules/flake8_comprehensions/fixes.rs
@@ -961,8 +961,30 @@ pub fn fix_unnecessary_comprehension(
                 whitespace_before_args: ParenthesizableWhitespace::default(),
             }));
         }
+        Expression::DictComp(inner) => {
+            body.value = Expression::Call(Box::new(Call {
+                func: Box::new(Expression::Name(Box::new(Name {
+                    value: "dict",
+                    lpar: vec![],
+                    rpar: vec![],
+                }))),
+                args: vec![Arg {
+                    value: inner.for_in.iter.clone(),
+                    keyword: None,
+                    equal: None,
+                    comma: None,
+                    star: "",
+                    whitespace_after_star: ParenthesizableWhitespace::default(),
+                    whitespace_after_arg: ParenthesizableWhitespace::default(),
+                }],
+                lpar: vec![],
+                rpar: vec![],
+                whitespace_after_func: ParenthesizableWhitespace::default(),
+                whitespace_before_args: ParenthesizableWhitespace::default(),
+            }));
+        }
         _ => {
-            bail!("Expected Expression::ListComp | Expression:SetComp");
+            bail!("Expected Expression::ListComp | Expression:SetComp | Expression:DictComp");
         }
     }
 

diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/helpers.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/helpers.rs
@@ -1,6 +1,6 @@
 use rustpython_parser::ast::{Expr, ExprKind, Keyword};
 
-pub fn function_name(func: &Expr) -> Option<&str> {
+pub fn expr_name(func: &Expr) -> Option<&str> {
     if let ExprKind::Name { id, .. } = &func.node {
         Some(id)
     } else {
@@ -20,7 +20,7 @@ pub fn exactly_one_argument_with_matching_function<'a>(
     if args.len() != 1 {
         return None;
     }
-    if function_name(func)? != name {
+    if expr_name(func)? != name {
         return None;
     }
     Some(&args[0].node)
@@ -31,7 +31,7 @@ pub fn first_argument_with_matching_function<'a>(
     func: &Expr,
     args: &'a [Expr],
 ) -> Option<&'a ExprKind> {
-    if function_name(func)? == name {
+    if expr_name(func)? == name {
         Some(&args.first()?.node)
     } else {
         None

diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/mod.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/mod.rs
@@ -2,7 +2,9 @@ pub use unnecessary_call_around_sorted::{
     unnecessary_call_around_sorted, UnnecessaryCallAroundSorted,
 };
 pub use unnecessary_collection_call::{unnecessary_collection_call, UnnecessaryCollectionCall};
-pub use unnecessary_comprehension::{unnecessary_comprehension, UnnecessaryComprehension};
+pub use unnecessary_comprehension::{
+    unnecessary_dict_comprehension, unnecessary_list_set_comprehension, UnnecessaryComprehension,
+};
 pub use unnecessary_double_cast_or_process::{
     unnecessary_double_cast_or_process, UnnecessaryDoubleCastOrProcess,
 };

diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_call_around_sorted.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_call_around_sorted.rs
@@ -59,7 +59,7 @@ pub fn unnecessary_call_around_sorted(
     func: &Expr,
     args: &[Expr],
 ) {
-    let Some(outer) = helpers::function_name(func) else {
+    let Some(outer) = helpers::expr_name(func) else {
         return;
     };
     if !(outer == "list" || outer == "reversed") {
@@ -71,7 +71,7 @@ pub fn unnecessary_call_around_sorted(
     let ExprKind::Call { func, .. } = &arg.node else {
         return;
     };
-    let Some(inner) = helpers::function_name(func) else {
+    let Some(inner) = helpers::expr_name(func) else {
         return;
     };
     if inner != "sorted" {

diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_collection_call.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_collection_call.rs
@@ -41,7 +41,7 @@ pub fn unnecessary_collection_call(
     if !args.is_empty() {
         return;
     }
-    let Some(id) = helpers::function_name(func) else {
+    let Some(id) = helpers::expr_name(func) else {
         return;
     };
     match id {

diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_comprehension.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_comprehension.rs
@@ -11,6 +11,27 @@ use crate::rules::flake8_comprehensions::fixes;
 
 use super::helpers;
 
+/// ## What it does
+/// Checks for unnecessary `dict`, `list`, and `set` comprehension.
+///
+/// ## Why is this bad?
+/// It's unnecessary to use a `dict`/`list`/`set` comprehension to build a
+/// data structure if the elements are unchanged. Wrap the iterable with
+/// `dict()`, `list()`, or `set()` instead.
+///
+/// ## Examples
+/// ```python
+/// {a: b for a, b in iterable}
+/// [x for x in iterable]
+/// {x for x in iterable}
+/// ```
+///
+/// Use instead:
+/// ```python
+/// dict(iterable)
+/// list(iterable)
+/// set(iterable)
+/// ```
 #[violation]
 pub struct UnnecessaryComprehension {
     pub obj_type: String,
@@ -29,33 +50,12 @@ impl AlwaysAutofixableViolation for UnnecessaryComprehension {
     }
 }
 
-/// C416
-pub fn unnecessary_comprehension(
-    checker: &mut Checker,
-    expr: &Expr,
-    elt: &Expr,
-    generators: &[Comprehension],
-) {
-    if generators.len() != 1 {
-        return;
-    }
-    let generator = &generators[0];
-    if !(generator.ifs.is_empty() && generator.is_async == 0) {
-        return;
-    }
-    let Some(elt_id) = helpers::function_name(elt) else {
-        return;
-    };
-
-    let Some(target_id) = helpers::function_name(&generator.target) else {
-        return;
-    };
-    if elt_id != target_id {
-        return;
-    }
+/// Add diagnostic for C416 based on the expression node id.
+fn add_diagnostic(checker: &mut Checker, expr: &Expr) {
     let id = match &expr.node {
         ExprKind::ListComp { .. } => "list",
         ExprKind::SetComp { .. } => "set",
+        ExprKind::DictComp { .. } => "dict",
         _ => return,
     };
     if !checker.ctx.is_builtin(id) {
@@ -77,3 +77,71 @@ pub fn unnecessary_comprehension(
     }
     checker.diagnostics.push(diagnostic);
 }
+
+/// C416
+pub fn unnecessary_dict_comprehension(
+    checker: &mut Checker,
+    expr: &Expr,
+    key: &Expr,
+    value: &Expr,
+    generators: &[Comprehension],
+) {
+    if generators.len() != 1 {
+        return;
+    }
+    let generator = &generators[0];
+    if !(generator.ifs.is_empty() && generator.is_async == 0) {
+        return;
+    }
+    let Some(key_id) = helpers::expr_name(key) else {
+        return;
+    };
+    let Some(value_id) = helpers::expr_name(value) else {
+        return;
+    };
+    let ExprKind::Tuple { elts, .. } = &generator.target.node else {
+        return;
+    };
+    if elts.len() != 2 {
+        return;
+    }
+    let Some(target_key_id) = helpers::expr_name(&elts[0]) else {
+        return;
+    };
+    if target_key_id != key_id {
+        return;
+    }
+    let Some(target_value_id) = helpers::expr_name(&elts[1]) else {
+        return;
+    };
+    if target_value_id != value_id {
+        return;
+    }
+    add_diagnostic(checker, expr);
+}
+
+/// C416
+pub fn unnecessary_list_set_comprehension(
+    checker: &mut Checker,
+    expr: &Expr,
+    elt: &Expr,
+    generators: &[Comprehension],
+) {
+    if generators.len() != 1 {
+        return;
+    }
+    let generator = &generators[0];
+    if !(generator.ifs.is_empty() && generator.is_async == 0) {
+        return;
+    }
+    let Some(elt_id) = helpers::expr_name(elt) else {
+        return;
+    };
+    let Some(target_id) = helpers::expr_name(&generator.target) else {
+        return;
+    };
+    if elt_id != target_id {
+        return;
+    }
+    add_diagnostic(checker, expr);
+}
diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_double_cast_or_process.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_double_cast_or_process.rs
@@ -80,7 +80,7 @@ pub fn unnecessary_double_cast_or_process(
         )
     }
 
-    let Some(outer) = helpers::function_name(func) else {
+    let Some(outer) = helpers::expr_name(func) else {
         return;
     };
     if !(outer == "list"
@@ -97,7 +97,7 @@ pub fn unnecessary_double_cast_or_process(
     let ExprKind::Call { func, .. } = &arg.node else {
         return;
     };
-    let Some(inner) = helpers::function_name(func) else {
+    let Some(inner) = helpers::expr_name(func) else {
         return;
     };
     if !checker.ctx.is_builtin(inner) || !checker.ctx.is_builtin(outer) {

diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_map.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_map.rs
@@ -84,7 +84,7 @@ pub fn unnecessary_map(
         )
     }
 
-    let Some(id) = helpers::function_name(func)  else {
+    let Some(id) = helpers::expr_name(func)  else {
         return;
     };
     match id {
@@ -96,7 +96,7 @@ pub fn unnecessary_map(
             // Exclude the parent if already matched by other arms
             if let Some(parent) = parent {
                 if let ExprKind::Call { func: f, .. } = &parent.node {
-                    if let Some(id_parent) = helpers::function_name(f) {
+                    if let Some(id_parent) = helpers::expr_name(f) {
                         if id_parent == "dict" || id_parent == "set" || id_parent == "list" {
                             return;
                         }

diff --git a/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_subscript_reversal.rs b/crates/ruff/src/rules/flake8_comprehensions/rules/unnecessary_subscript_reversal.rs
@@ -32,7 +32,7 @@ pub fn unnecessary_subscript_reversal(
     let Some(first_arg) = args.first() else {
         return;
     };
-    let Some(id) = helpers::function_name(func) else {
+    let Some(id) = helpers::expr_name(func) else {
         return;
     };
     if !(id == "set" || id == "sorted" || id == "reversed") {