Restyle flake8_comprehensions::check to reduce indent (#462)

src/flake8_comprehensions/checks.rs

@@ -4,57 +4,75 @@ use rustpython_ast::{Comprehension, Constant, Expr, ExprKind, KeywordData, Locat
 use crate::ast::types::Range;
 use crate::checks::{Check, CheckKind};
 
+fn function_name(func: &Expr) -> Option<&str> {
+    if let ExprKind::Name { id, .. } = &func.node {
+        Some(id)
+    } else {
+        None
+    }
+}
+
+fn exactly_one_argument_with_matching_function<'a>(
+    name: &str,
+    func: &Expr,
+    args: &'a [Expr],
+) -> Option<&'a ExprKind> {
+    if args.len() != 1 {
+        return None;
+    }
+    if function_name(func)? != name {
+        return None;
+    }
+    Some(&args[0].node)
+}
+
+fn first_argument_with_matching_function<'a>(
+    name: &str,
+    func: &Expr,
+    args: &'a [Expr],
+) -> Option<&'a ExprKind> {
+    if function_name(func)? != name {
+        return None;
+    }
+    Some(&args.first()?.node)
+}
+
 /// Check `list(generator)` compliance.
 pub fn unnecessary_generator_list(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if args.len() == 1 {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "list" {
-                if let ExprKind::GeneratorExp { .. } = &args[0].node {
-                    return Some(Check::new(
-                        CheckKind::UnnecessaryGeneratorList,
-                        Range::from_located(expr),
-                    ));
-                }
-            }
-        }
+    let argument = exactly_one_argument_with_matching_function("list", func, args)?;
+    if let ExprKind::GeneratorExp { .. } = argument {
+        return Some(Check::new(
+            CheckKind::UnnecessaryGeneratorList,
+            Range::from_located(expr),
+        ));
     }
     None
 }
 
 /// Check `set(generator)` compliance.
 pub fn unnecessary_generator_set(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if args.len() == 1 {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "set" {
-                if let ExprKind::GeneratorExp { .. } = &args[0].node {
-                    return Some(Check::new(
-                        CheckKind::UnnecessaryGeneratorSet,
-                        Range::from_located(expr),
-                    ));
-                }
-            }
-        }
+    let argument = exactly_one_argument_with_matching_function("set", func, args)?;
+    if let ExprKind::GeneratorExp { .. } = argument {
+        return Some(Check::new(
+            CheckKind::UnnecessaryGeneratorSet,
+            Range::from_located(expr),
+        ));
     }
     None
 }
 
 /// Check `dict((x, y) for x, y in iterable)` compliance.
 pub fn unnecessary_generator_dict(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if args.len() == 1 {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "dict" {
-                if let ExprKind::GeneratorExp { elt, .. } = &args[0].node {
-                    match &elt.node {
-                        ExprKind::Tuple { elts, .. } if elts.len() == 2 => {
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryGeneratorDict,
-                                Range::from_located(expr),
-                            ));
-                        }
-                        _ => {}
-                    }
-                }
+    let argument = exactly_one_argument_with_matching_function("dict", func, args)?;
+    if let ExprKind::GeneratorExp { elt, .. } = argument {
+        match &elt.node {
+            ExprKind::Tuple { elts, .. } if elts.len() == 2 => {
+                return Some(Check::new(
+                    CheckKind::UnnecessaryGeneratorDict,
+                    Range::from_located(expr),
+                ));
             }
+            _ => {}
         }
     }
     None
@@ -66,17 +84,12 @@ pub fn unnecessary_list_comprehension_set(
     func: &Expr,
     args: &[Expr],
 ) -> Option<Check> {
-    if args.len() == 1 {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "set" {
-                if let ExprKind::ListComp { .. } = &args[0].node {
-                    return Some(Check::new(
-                        CheckKind::UnnecessaryListComprehensionSet,
-                        Range::from_located(expr),
-                    ));
-                }
-            }
-        }
+    let argument = exactly_one_argument_with_matching_function("set", func, args)?;
+    if let ExprKind::ListComp { .. } = &argument {
+        return Some(Check::new(
+            CheckKind::UnnecessaryListComprehensionSet,
+            Range::from_located(expr),
+        ));
     }
     None
 }
@@ -87,21 +100,16 @@ pub fn unnecessary_list_comprehension_dict(
     func: &Expr,
     args: &[Expr],
 ) -> Option<Check> {
-    if args.len() == 1 {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "dict" {
-                if let ExprKind::ListComp { elt, .. } = &args[0].node {
-                    match &elt.node {
-                        ExprKind::Tuple { elts, .. } if elts.len() == 2 => {
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryListComprehensionDict,
-                                Range::from_located(expr),
-                            ));
-                        }
-                        _ => {}
-                    }
-                }
+    let argument = exactly_one_argument_with_matching_function("dict", func, args)?;
+    if let ExprKind::ListComp { elt, .. } = &argument {
+        match &elt.node {
+            ExprKind::Tuple { elts, .. } if elts.len() == 2 => {
+                return Some(Check::new(
+                    CheckKind::UnnecessaryListComprehensionDict,
+                    Range::from_located(expr),
+                ));
             }
+            _ => {}
         }
     }
     None
@@ -109,82 +117,38 @@ pub fn unnecessary_list_comprehension_dict(
 
 /// Check `set([1, 2])` compliance.
 pub fn unnecessary_literal_set(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if args.len() == 1 {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "set" {
-                match &args[0].node {
-                    ExprKind::List { .. } => {
-                        return Some(Check::new(
-                            CheckKind::UnnecessaryLiteralSet("list".to_string()),
-                            Range::from_located(expr),
-                        ));
-                    }
-                    ExprKind::Tuple { .. } => {
-                        return Some(Check::new(
-                            CheckKind::UnnecessaryLiteralSet("tuple".to_string()),
-                            Range::from_located(expr),
-                        ));
-                    }
-                    _ => {}
-                }
-            }
-        }
-    }
-    None
+    let argument = exactly_one_argument_with_matching_function("set", func, args)?;
+    let kind = match argument {
+        ExprKind::List { .. } => "list",
+        ExprKind::Tuple { .. } => "tuple",
+        _ => return None,
+    };
+    Some(Check::new(
+        CheckKind::UnnecessaryLiteralSet(kind.to_string()),
+        Range::from_located(expr),
+    ))
 }
 
 /// Check `dict([(1, 2)])` compliance.
 pub fn unnecessary_literal_dict(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if args.len() == 1 {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "dict" {
-                match &args[0].node {
-                    ExprKind::Tuple { elts, .. } => {
-                        if let Some(elt) = elts.first() {
-                            match &elt.node {
-                                // dict((1, 2), ...))
-                                ExprKind::Tuple { elts, .. } if elts.len() == 2 => {
-                                    return Some(Check::new(
-                                        CheckKind::UnnecessaryLiteralDict("tuple".to_string()),
-                                        Range::from_located(expr),
-                                    ));
-                                }
-                                _ => {}
-                            }
-                        } else {
-                            // dict(())
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryLiteralDict("tuple".to_string()),
-                                Range::from_located(expr),
-                            ));
-                        }
-                    }
-                    ExprKind::List { elts, .. } => {
-                        if let Some(elt) = elts.first() {
-                            match &elt.node {
-                                // dict([(1, 2), ...])
-                                ExprKind::Tuple { elts, .. } if elts.len() == 2 => {
-                                    return Some(Check::new(
-                                        CheckKind::UnnecessaryLiteralDict("list".to_string()),
-                                        Range::from_located(expr),
-                                    ));
-                                }
-                                _ => {}
-                            }
-                        } else {
-                            // dict([])
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryLiteralDict("list".to_string()),
-                                Range::from_located(expr),
-                            ));
-                        }
-                    }
-                    _ => {}
-                }
-            }
+    let argument = exactly_one_argument_with_matching_function("dict", func, args)?;
+    let (kind, elts) = match argument {
+        ExprKind::Tuple { elts, .. } => ("tuple", elts),
+        ExprKind::List { elts, .. } => ("list", elts),
+        _ => return None,
+    };
+
+    if let Some(elt) = elts.first() {
+        // dict((1, 2), ...)) or dict([(1, 2), ...])
+        if !matches!(&elt.node, ExprKind::Tuple { elts, .. } if elts.len() == 2) {
+            return None;
         }
     }
-    None
+
+    Some(Check::new(
+        CheckKind::UnnecessaryLiteralDict(kind.to_string()),
+        Range::from_located(expr),
+    ))
 }
 
 pub fn unnecessary_collection_call(
@@ -193,26 +157,21 @@ pub fn unnecessary_collection_call(
     args: &[Expr],
     keywords: &[Located<KeywordData>],
 ) -> Option<Check> {
-    if args.is_empty() {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "list" || id == "tuple" {
-                // list() or tuple()
-                return Some(Check::new(
-                    CheckKind::UnnecessaryCollectionCall(id.to_string()),
-                    Range::from_located(expr),
-                ));
-            } else if id == "dict" {
-                // dict() or dict(a=1)
-                if keywords.is_empty() || keywords.iter().all(|kw| kw.node.arg.is_some()) {
-                    return Some(Check::new(
-                        CheckKind::UnnecessaryCollectionCall(id.to_string()),
-                        Range::from_located(expr),
-                    ));
-                }
-            }
-        }
+    if !args.is_empty() {
+        return None;
     }
-    None
+    let id = function_name(func)?;
+    match id {
+        "list" | "tuple" => {
+            // list() or tuple()
+        }
+        "dict" if keywords.is_empty() || keywords.iter().all(|kw| kw.node.arg.is_some()) => (),
+        _ => return None,
+    };
+    Some(Check::new(
+        CheckKind::UnnecessaryCollectionCall(id.to_string()),
+        Range::from_located(expr),
+    ))
 }
 
 pub fn unnecessary_literal_within_tuple_call(
@@ -220,29 +179,16 @@ pub fn unnecessary_literal_within_tuple_call(
     func: &Expr,
     args: &[Expr],
 ) -> Option<Check> {
-    if let ExprKind::Name { id, .. } = &func.node {
-        if id == "tuple" {
-            if let Some(arg) = args.first() {
-                match &arg.node {
-                    ExprKind::Tuple { .. } => {
-                        return Some(Check::new(
-                            CheckKind::UnnecessaryLiteralWithinTupleCall("tuple".to_string()),
-                            Range::from_located(expr),
-                        ));
-                    }
-                    ExprKind::List { .. } => {
-                        return Some(Check::new(
-                            CheckKind::UnnecessaryLiteralWithinTupleCall("list".to_string()),
-                            Range::from_located(expr),
-                        ));
-                    }
-                    _ => {}
-                }
-            }
-        }
-    }
-
-    None
+    let argument = first_argument_with_matching_function("tuple", func, args)?;
+    let argument_kind = match argument {
+        ExprKind::Tuple { .. } => "tuple",
+        ExprKind::List { .. } => "list",
+        _ => return None,
+    };
+    Some(Check::new(
+        CheckKind::UnnecessaryLiteralWithinTupleCall(argument_kind.to_string()),
+        Range::from_located(expr),
+    ))
 }
 
 pub fn unnecessary_literal_within_list_call(
@@ -250,62 +196,40 @@ pub fn unnecessary_literal_within_list_call(
     func: &Expr,
     args: &[Expr],
 ) -> Option<Check> {
-    if let ExprKind::Name { id, .. } = &func.node {
-        if id == "list" {
-            if let Some(arg) = args.first() {
-                match &arg.node {
-                    ExprKind::Tuple { .. } => {
-                        return Some(Check::new(
-                            CheckKind::UnnecessaryLiteralWithinListCall("tuple".to_string()),
-                            Range::from_located(expr),
-                        ));
-                    }
-                    ExprKind::List { .. } => {
-                        return Some(Check::new(
-                            CheckKind::UnnecessaryLiteralWithinListCall("list".to_string()),
-                            Range::from_located(expr),
-                        ));
-                    }
-                    _ => {}
-                }
-            }
-        }
-    }
-
-    None
+    let argument = first_argument_with_matching_function("list", func, args)?;
+    let argument_kind = match argument {
+        ExprKind::Tuple { .. } => "tuple",
+        ExprKind::List { .. } => "list",
+        _ => return None,
+    };
+    Some(Check::new(
+        CheckKind::UnnecessaryLiteralWithinListCall(argument_kind.to_string()),
+        Range::from_located(expr),
+    ))
 }
 
 pub fn unnecessary_list_call(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if let ExprKind::Name { id, .. } = &func.node {
-        if id == "list" {
-            if let Some(arg) = args.first() {
-                if let ExprKind::ListComp { .. } = &arg.node {
-                    return Some(Check::new(
-                        CheckKind::UnnecessaryListCall,
-                        Range::from_located(expr),
-                    ));
-                }
-            }
-        }
+    let argument = first_argument_with_matching_function("list", func, args)?;
+    if let ExprKind::ListComp { .. } = argument {
+        return Some(Check::new(
+            CheckKind::UnnecessaryListCall,
+            Range::from_located(expr),
+        ));
     }
     None
 }
 
 pub fn unnecessary_call_around_sorted(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if let ExprKind::Name { id: outer, .. } = &func.node {
-        if outer == "list" || outer == "reversed" {
-            if let Some(arg) = args.first() {
-                if let ExprKind::Call { func, .. } = &arg.node {
-                    if let ExprKind::Name { id: inner, .. } = &func.node {
-                        if inner == "sorted" {
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryCallAroundSorted(outer.to_string()),
-                                Range::from_located(expr),
-                            ));
-                        }
-                    }
-                }
-            }
+    let outer = function_name(func)?;
+    if !(outer == "list" || outer == "reversed") {
+        return None;
+    }
+    if let ExprKind::Call { func, .. } = &args.first()?.node {
+        if function_name(func)? == "sorted" {
+            return Some(Check::new(
+                CheckKind::UnnecessaryCallAroundSorted(outer.to_string()),
+                Range::from_located(expr),
+            ));
         }
     }
     None
@@ -316,91 +240,65 @@ pub fn unnecessary_double_cast_or_process(
     func: &Expr,
     args: &[Expr],
 ) -> Option<Check> {
-    if let ExprKind::Name { id: outer, .. } = &func.node {
-        if outer == "list"
-            || outer == "tuple"
-            || outer == "set"
-            || outer == "reversed"
-            || outer == "sorted"
+    let outer = function_name(func)?;
+    if !["list", "tuple", "set", "reversed", "sorted"].contains(&outer) {
+        return None;
+    }
+
+    fn new_check(inner: &str, outer: &str, expr: &Expr) -> Check {
+        Check::new(
+            CheckKind::UnnecessaryDoubleCastOrProcess(inner.to_string(), outer.to_string()),
+            Range::from_located(expr),
+        )
+    }
+
+    if let ExprKind::Call { func, .. } = &args.first()?.node {
+        let inner = function_name(func)?;
+        // Ex) set(tuple(...))
+        if (outer == "set" || outer == "sorted")
+            && (inner == "list" || inner == "tuple" || inner == "reversed" || inner == "sorted")
         {
-            if let Some(arg) = args.first() {
-                if let ExprKind::Call { func, .. } = &arg.node {
-                    if let ExprKind::Name { id: inner, .. } = &func.node {
-                        // Ex) set(tuple(...))
-                        if (outer == "set" || outer == "sorted")
-                            && (inner == "list"
-                                || inner == "tuple"
-                                || inner == "reversed"
-                                || inner == "sorted")
-                        {
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryDoubleCastOrProcess(
-                                    inner.to_string(),
-                                    outer.to_string(),
-                                ),
-                                Range::from_located(expr),
-                            ));
-                        }
+            return Some(new_check(inner, outer, expr));
+        }
 
-                        // Ex) list(tuple(...))
-                        if (outer == "list" || outer == "tuple")
-                            && (inner == "list" || inner == "tuple")
-                        {
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryDoubleCastOrProcess(
-                                    inner.to_string(),
-                                    outer.to_string(),
-                                ),
-                                Range::from_located(expr),
-                            ));
-                        }
+        // Ex) list(tuple(...))
+        if (outer == "list" || outer == "tuple") && (inner == "list" || inner == "tuple") {
+            return Some(new_check(inner, outer, expr));
+        }
 
-                        // Ex) set(set(...))
-                        if outer == "set" && inner == "set" {
-                            return Some(Check::new(
-                                CheckKind::UnnecessaryDoubleCastOrProcess(
-                                    inner.to_string(),
-                                    outer.to_string(),
-                                ),
-                                Range::from_located(expr),
-                            ));
-                        }
-                    }
-                }
-            }
+        // Ex) set(set(...))
+        if outer == "set" && inner == "set" {
+            return Some(new_check(inner, outer, expr));
         }
     }
     None
 }
 
 pub fn unnecessary_subscript_reversal(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if let Some(first_arg) = args.first() {
-        if let ExprKind::Name { id, .. } = &func.node {
-            if id == "set" || id == "sorted" || id == "reversed" {
-                if let ExprKind::Subscript { slice, .. } = &first_arg.node {
-                    if let ExprKind::Slice { lower, upper, step } = &slice.node {
-                        if lower.is_none() && upper.is_none() {
-                            if let Some(step) = step {
-                                if let ExprKind::UnaryOp {
-                                    op: Unaryop::USub,
-                                    operand,
-                                } = &step.node
-                                {
-                                    if let ExprKind::Constant {
-                                        value: Constant::Int(val),
-                                        ..
-                                    } = &operand.node
-                                    {
-                                        if *val == BigInt::from(1) {
-                                            return Some(Check::new(
-                                                CheckKind::UnnecessarySubscriptReversal(
-                                                    id.to_string(),
-                                                ),
-                                                Range::from_located(expr),
-                                            ));
-                                        }
-                                    }
-                                }
+    let first_arg = args.first()?;
+    let id = function_name(func)?;
+    if !["set", "sorted", "reversed"].contains(&id) {
+        return None;
+    }
+    if let ExprKind::Subscript { slice, .. } = &first_arg.node {
+        if let ExprKind::Slice { lower, upper, step } = &slice.node {
+            if lower.is_none() && upper.is_none() {
+                if let Some(step) = step {
+                    if let ExprKind::UnaryOp {
+                        op: Unaryop::USub,
+                        operand,
+                    } = &step.node
+                    {
+                        if let ExprKind::Constant {
+                            value: Constant::Int(val),
+                            ..
+                        } = &operand.node
+                        {
+                            if *val == BigInt::from(1) {
+                                return Some(Check::new(
+                                    CheckKind::UnnecessarySubscriptReversal(id.to_string()),
+                                    Range::from_located(expr),
+                                ));
                             }
                         }
                     }
@@ -416,90 +314,65 @@ pub fn unnecessary_comprehension(
     elt: &Expr,
     generators: &[Comprehension],
 ) -> Option<Check> {
-    if generators.len() == 1 {
-        let generator = &generators[0];
-        if generator.ifs.is_empty() && generator.is_async == 0 {
-            if let ExprKind::Name { id: elt_id, .. } = &elt.node {
-                if let ExprKind::Name { id: target_id, .. } = &generator.target.node {
-                    if elt_id == target_id {
-                        match &expr.node {
-                            ExprKind::ListComp { .. } => {
-                                return Some(Check::new(
-                                    CheckKind::UnnecessaryComprehension("list".to_string()),
-                                    Range::from_located(expr),
-                                ))
-                            }
-                            ExprKind::SetComp { .. } => {
-                                return Some(Check::new(
-                                    CheckKind::UnnecessaryComprehension("set".to_string()),
-                                    Range::from_located(expr),
-                                ))
-                            }
-                            _ => {}
-                        };
-                    }
-                }
-            }
-        }
+    if generators.len() != 1 {
+        return None;
     }
-
-    None
+    let generator = &generators[0];
+    if !(generator.ifs.is_empty() && generator.is_async == 0) {
+        return None;
+    }
+    let elt_id = function_name(elt)?;
+    let target_id = function_name(&generator.target)?;
+    if elt_id != target_id {
+        return None;
+    }
+    let expr_kind = match &expr.node {
+        ExprKind::ListComp { .. } => "list",
+        ExprKind::SetComp { .. } => "set",
+        _ => return None,
+    };
+    Some(Check::new(
+        CheckKind::UnnecessaryComprehension(expr_kind.to_string()),
+        Range::from_located(expr),
+    ))
 }
 
 pub fn unnecessary_map(expr: &Expr, func: &Expr, args: &[Expr]) -> Option<Check> {
-    if let ExprKind::Name { id, .. } = &func.node {
-        if id == "map" {
-            if args.len() == 2 {
-                if let ExprKind::Lambda { .. } = &args[0].node {
-                    return Some(Check::new(
-                        CheckKind::UnnecessaryMap("generator".to_string()),
-                        Range::from_located(expr),
-                    ));
+    fn new_check(kind: &str, expr: &Expr) -> Check {
+        Check::new(
+            CheckKind::UnnecessaryMap(kind.to_string()),
+            Range::from_located(expr),
+        )
+    }
+    let id = function_name(func)?;
+    match id {
+        "map" => {
+            if args.len() == 2 && matches!(&args[0].node, ExprKind::Lambda { .. }) {
+                return Some(new_check("generator", expr));
+            }
+        }
+        "list" | "set" => {
+            if let ExprKind::Call { func, args, .. } = &args.first()?.node {
+                let argument = first_argument_with_matching_function("map", func, args)?;
+                if let ExprKind::Lambda { .. } = argument {
+                    return Some(new_check(id, expr));
                 }
             }
-        } else if id == "list" || id == "set" {
-            if let Some(arg) = args.first() {
-                if let ExprKind::Call { func, args, .. } = &arg.node {
-                    if let ExprKind::Name { id: f, .. } = &func.node {
-                        if f == "map" {
-                            if let Some(arg) = args.first() {
-                                if let ExprKind::Lambda { .. } = &arg.node {
-                                    return Some(Check::new(
-                                        CheckKind::UnnecessaryMap(id.to_string()),
-                                        Range::from_located(expr),
-                                    ));
-                                }
-                            }
-                        }
-                    }
-                }
-            }
-        } else if id == "dict" {
+        }
+        "dict" => {
             if args.len() == 1 {
                 if let ExprKind::Call { func, args, .. } = &args[0].node {
-                    if let ExprKind::Name { id: f, .. } = &func.node {
-                        if f == "map" {
-                            if let Some(arg) = args.first() {
-                                if let ExprKind::Lambda { body, .. } = &arg.node {
-                                    match &body.node {
-                                        ExprKind::Tuple { elts, .. }
-                                        | ExprKind::List { elts, .. }
-                                            if elts.len() == 2 =>
-                                        {
-                                            return Some(Check::new(
-                                                CheckKind::UnnecessaryMap(id.to_string()),
-                                                Range::from_located(expr),
-                                            ))
-                                        }
-                                        _ => {}
-                                    }
-                                }
-                            }
+                    let argument = first_argument_with_matching_function("map", func, args)?;
+                    if let ExprKind::Lambda { body, .. } = &argument {
+                        if matches!(&body.node, ExprKind::Tuple { elts, .. } | ExprKind::List { elts, .. } if elts.len() == 2)
+                        {
+                            return Some(new_check(id, expr));
                         }
                     }
                 }
             }
         }
+        _ => (),
     }
     None
 }