DRY up utilities in flake8_comprehensions/fixes.rs (#556)

src/cst/matchers.rs

@@ -0,0 +1,19 @@
+use anyhow::Result;
+use libcst_native::Module;
+use rustpython_ast::Located;
+
+use crate::ast::types::Range;
+use crate::source_code_locator::SourceCodeLocator;
+
+pub fn match_tree<'a, T>(
+    locator: &'a SourceCodeLocator,
+    located: &'a Located<T>,
+) -> Result<Module<'a>> {
+    match libcst_native::parse_module(
+        locator.slice_source_code_range(&Range::from_located(located)),
+        None,
+    ) {
+        Ok(module) => Ok(module),
+        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
+    }
+}

src/cst/mod.rs

@@ -1 +1,2 @@
 pub mod helpers;
+pub mod matchers;

src/flake8_comprehensions/fixes.rs

@@ -1,54 +1,57 @@
 use anyhow::Result;
 use libcst_native::{
-    Arg, Codegen, Dict, DictComp, Element, Expression, LeftCurlyBrace, LeftParen,
+    Arg, Call, Codegen, Dict, DictComp, Element, Expr, Expression, LeftCurlyBrace, LeftParen,
-    LeftSquareBracket, List, ListComp, ParenthesizableWhitespace, RightCurlyBrace, RightParen,
+    LeftSquareBracket, List, ListComp, Module, ParenthesizableWhitespace, RightCurlyBrace,
-    RightSquareBracket, Set, SetComp, SimpleWhitespace, SmallStatement, Statement, Tuple,
+    RightParen, RightSquareBracket, Set, SetComp, SimpleWhitespace, SmallStatement, Statement,
+    Tuple,
 };
-use rustpython_ast::Expr;
 
-use crate::ast::types::Range;
 use crate::autofix::Fix;
+use crate::cst::matchers::match_tree;
 use crate::source_code_locator::SourceCodeLocator;
 
-/// (C400) Convert `list(x for x in y)` to `[x for x in y]`.
+fn match_expr<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Expr<'b>> {
-pub fn fix_unnecessary_generator_list(locator: &SourceCodeLocator, expr: &Expr) -> Result<Fix> {
+    if let Some(Statement::Simple(expr)) = module.body.first_mut() {
-    // Module(SimpleStatementLine(Expr(Call(GeneratorExp)))) ->
+        if let Some(SmallStatement::Expr(expr)) = expr.body.first_mut() {
-    // Module(SimpleStatementLine(Expr(ListComp)))
+            Ok(expr)
-    let mut tree = match libcst_native::parse_module(
-        locator.slice_source_code_range(&Range::from_located(expr)),
-        None,
-    ) {
-        Ok(m) => m,
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
-    };
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
-        body
         } else {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
+            Err(anyhow::anyhow!(
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
                 "Expected node to be: SmallStatement::Expr."
-        ));
+            ))
-    };
+        }
-    let call = if let Expression::Call(call) = &body.value {
-        call
     } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
+        Err(anyhow::anyhow!("Expected node to be: Statement::Simple."))
-    };
+    }
-    let (arg, whitespace_after_arg) = if let Some(Arg {
+}
-        value,
+
-        whitespace_after_arg,
+fn match_call<'a, 'b>(expr: &'a mut Expr<'b>) -> Result<&'a mut Call<'b>> {
-        ..
+    if let Expression::Call(call) = &mut expr.value {
-    }) = call.args.first()
+        Ok(call)
-    {
-        (value, whitespace_after_arg)
     } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
+        Err(anyhow::anyhow!("Expected node to be: Expression::Call."))
-    };
+    }
-    let generator_exp = if let Expression::GeneratorExp(generator_exp) = &arg {
+}
+
+fn match_arg<'a, 'b>(call: &'a Call<'b>) -> Result<&'a Arg<'b>> {
+    if let Some(arg) = call.args.first() {
+        Ok(arg)
+    } else {
+        Err(anyhow::anyhow!("Expected node to be: Arg."))
+    }
+}
+
+/// (C400) Convert `list(x for x in y)` to `[x for x in y]`.
+pub fn fix_unnecessary_generator_list(
+    locator: &SourceCodeLocator,
+    expr: &rustpython_ast::Expr,
+) -> Result<Fix> {
+    // Expr(Call(GeneratorExp)))) -> Expr(ListComp)))
+    let mut tree = match_tree(locator, expr)?;
+    let mut body = match_expr(&mut tree)?;
+    let call = match_call(body)?;
+    let arg = match_arg(call)?;
+
+    let generator_exp = if let Expression::GeneratorExp(generator_exp) = &arg.value {
         generator_exp
     } else {
         return Err(anyhow::anyhow!(
@@ -63,7 +66,7 @@ pub fn fix_unnecessary_generator_list(locator: &SourceCodeLocator, expr: &Expr)
             whitespace_after: call.whitespace_before_args.clone(),
         },
         rbracket: RightSquareBracket {
-            whitespace_before: whitespace_after_arg.clone(),
+            whitespace_before: arg.whitespace_after_arg.clone(),
         },
         lpar: generator_exp.lpar.clone(),
         rpar: generator_exp.rpar.clone(),
@@ -80,44 +83,17 @@ pub fn fix_unnecessary_generator_list(locator: &SourceCodeLocator, expr: &Expr)
 }
 
 /// (C401) Convert `set(x for x in y)` to `{x for x in y}`.
-pub fn fix_unnecessary_generator_set(locator: &SourceCodeLocator, expr: &Expr) -> Result<Fix> {
+pub fn fix_unnecessary_generator_set(
-    // Module(SimpleStatementLine(Expr(Call(GeneratorExp)))) ->
+    locator: &SourceCodeLocator,
-    // Module(SimpleStatementLine(Expr(SetComp)))
+    expr: &rustpython_ast::Expr,
-    let mut tree = match libcst_native::parse_module(
+) -> Result<Fix> {
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    // Expr(Call(GeneratorExp)))) -> Expr(SetComp)))
-        None,
+    let mut tree = match_tree(locator, expr)?;
-    ) {
+    let mut body = match_expr(&mut tree)?;
-        Ok(m) => m,
+    let call = match_call(body)?;
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
+    let arg = match_arg(call)?;
-    };
+
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
+    let generator_exp = if let Expression::GeneratorExp(generator_exp) = &arg.value {
-        body
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
-    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value,
-        whitespace_after_arg,
-        ..
-    }) = call.args.first()
-    {
-        (value, whitespace_after_arg)
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
-    };
-    let generator_exp = if let Expression::GeneratorExp(generator_exp) = &arg {
         generator_exp
     } else {
         return Err(anyhow::anyhow!(
@@ -132,7 +108,7 @@ pub fn fix_unnecessary_generator_set(locator: &SourceCodeLocator, expr: &Expr) -
             whitespace_after: call.whitespace_before_args.clone(),
         },
         rbrace: RightCurlyBrace {
-            whitespace_before: whitespace_after_arg.clone(),
+            whitespace_before: arg.whitespace_after_arg.clone(),
         },
         lpar: generator_exp.lpar.clone(),
         rpar: generator_exp.rpar.clone(),
@@ -150,42 +126,17 @@ pub fn fix_unnecessary_generator_set(locator: &SourceCodeLocator, expr: &Expr) -
 
 /// (C402) Convert `dict((x, x) for x in range(3))` to `{x: x for x in
 /// range(3)}`.
-pub fn fix_unnecessary_generator_dict(locator: &SourceCodeLocator, expr: &Expr) -> Result<Fix> {
+pub fn fix_unnecessary_generator_dict(
-    let mut tree = match libcst_native::parse_module(
+    locator: &SourceCodeLocator,
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    expr: &rustpython_ast::Expr,
-        None,
+) -> Result<Fix> {
-    ) {
+    let mut tree = match_tree(locator, expr)?;
-        Ok(m) => m,
+    let mut body = match_expr(&mut tree)?;
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
+    let call = match_call(body)?;
-    };
+    let arg = match_arg(call)?;
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
+
-        body
+    // Extract the (k, v) from `(k, v) for ...`.
-    } else {
+    let generator_exp = if let Expression::GeneratorExp(generator_exp) = &arg.value {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
-    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value,
-        whitespace_after_arg,
-        ..
-    }) = call.args.first()
-    {
-        (value, whitespace_after_arg)
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
-    };
-    let generator_exp = if let Expression::GeneratorExp(generator_exp) = &arg {
         generator_exp
     } else {
         return Err(anyhow::anyhow!(
@@ -220,7 +171,7 @@ pub fn fix_unnecessary_generator_dict(locator: &SourceCodeLocator, expr: &Expr)
             whitespace_after: call.whitespace_before_args.clone(),
         },
         rbrace: RightCurlyBrace {
-            whitespace_before: whitespace_after_arg.clone(),
+            whitespace_before: arg.whitespace_after_arg.clone(),
         },
         lpar: Default::default(),
         rpar: Default::default(),
@@ -241,45 +192,16 @@ pub fn fix_unnecessary_generator_dict(locator: &SourceCodeLocator, expr: &Expr)
 /// (C403) Convert `set([x for x in y])` to `{x for x in y}`.
 pub fn fix_unnecessary_list_comprehension_set(
     locator: &SourceCodeLocator,
-    expr: &Expr,
+    expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    // Module(SimpleStatementLine(Expr(Call(ListComp)))) ->
+    // Expr(Call(ListComp)))) ->
-    // Module(SimpleStatementLine(Expr(SetComp)))
+    // Expr(SetComp)))
-    let mut tree = match libcst_native::parse_module(
+    let mut tree = match_tree(locator, expr)?;
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    let mut body = match_expr(&mut tree)?;
-        None,
+    let call = match_call(body)?;
-    ) {
+    let arg = match_arg(call)?;
-        Ok(m) => m,
+
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
+    let list_comp = if let Expression::ListComp(list_comp) = &arg.value {
-    };
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
-    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value,
-        whitespace_after_arg,
-        ..
-    }) = call.args.first()
-    {
-        (value, whitespace_after_arg)
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
-    };
-    let list_comp = if let Expression::ListComp(list_comp) = arg {
         list_comp
     } else {
         return Err(anyhow::anyhow!(
@@ -294,7 +216,7 @@ pub fn fix_unnecessary_list_comprehension_set(
             whitespace_after: call.whitespace_before_args.clone(),
         },
         rbrace: RightCurlyBrace {
-            whitespace_before: whitespace_after_arg.clone(),
+            whitespace_before: arg.whitespace_after_arg.clone(),
         },
         lpar: list_comp.lpar.clone(),
         rpar: list_comp.rpar.clone(),
@@ -311,46 +233,19 @@ pub fn fix_unnecessary_list_comprehension_set(
 }
 
 /// (C405) Convert `set((1, 2))` to `{1, 2}`.
-pub fn fix_unnecessary_literal_set(locator: &SourceCodeLocator, expr: &Expr) -> Result<Fix> {
+pub fn fix_unnecessary_literal_set(
-    // Module(SimpleStatementLine(Expr(Call(List|Tuple)))) ->
+    locator: &SourceCodeLocator,
-    // Module(SimpleStatementLine(Expr(Set)))
+    expr: &rustpython_ast::Expr,
-    let mut tree = match libcst_native::parse_module(
+) -> Result<Fix> {
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    // Expr(Call(List|Tuple)))) -> Expr(Set)))
-        None,
+    let mut tree = match_tree(locator, expr)?;
-    ) {
+    let mut body = match_expr(&mut tree)?;
-        Ok(m) => m,
+    let mut call = match_call(body)?;
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
+    let arg = match_arg(call)?;
-    };
+
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
+    let elements = match &arg.value {
-        body
+        Expression::Tuple(inner) => &inner.elements,
-    } else {
+        Expression::List(inner) => &inner.elements,
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &mut body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
-    let (arg, whitespace_after_arg) = if let Some(Arg {
-        value,
-        whitespace_after_arg,
-        ..
-    }) = call.args.first_mut()
-    {
-        (value, whitespace_after_arg)
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
-    };
-    let elements = match arg {
-        Expression::Tuple(inner) => inner.elements.clone(),
-        Expression::List(inner) => inner.elements.clone(),
         _ => {
             return Err(anyhow::anyhow!(
                 "Expected node to be: Expression::Tuple | Expression::List."
@@ -362,12 +257,12 @@ pub fn fix_unnecessary_literal_set(locator: &SourceCodeLocator, expr: &Expr) ->
         call.args = vec![];
     } else {
         body.value = Expression::Set(Box::new(Set {
-            elements,
+            elements: elements.clone(),
             lbrace: LeftCurlyBrace {
                 whitespace_after: call.whitespace_before_args.clone(),
             },
             rbrace: RightCurlyBrace {
-                whitespace_before: whitespace_after_arg.clone(),
+                whitespace_before: arg.whitespace_after_arg.clone(),
             },
             lpar: Default::default(),
             rpar: Default::default(),
@@ -385,33 +280,14 @@ pub fn fix_unnecessary_literal_set(locator: &SourceCodeLocator, expr: &Expr) ->
 }
 
 /// (C408)
-pub fn fix_unnecessary_collection_call(locator: &SourceCodeLocator, expr: &Expr) -> Result<Fix> {
+pub fn fix_unnecessary_collection_call(
-    // Module(SimpleStatementLine(Expr(Call("list" | "tuple" | "dict")))) ->
+    locator: &SourceCodeLocator,
-    // Module(SimpleStatementLine(Expr(List|Tuple|Dict)))
+    expr: &rustpython_ast::Expr,
-    let mut tree = match libcst_native::parse_module(
+) -> Result<Fix> {
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    // Expr(Call("list" | "tuple" | "dict")))) -> Expr(List|Tuple|Dict)))
-        None,
+    let mut tree = match_tree(locator, expr)?;
-    ) {
+    let mut body = match_expr(&mut tree)?;
-        Ok(m) => m,
+    let call = match_call(body)?;
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
-    };
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
     let name = if let Expression::Name(name) = &call.func.as_ref() {
         name
     } else {
@@ -421,27 +297,27 @@ pub fn fix_unnecessary_collection_call(locator: &SourceCodeLocator, expr: &Expr)
     match name.value {
         "tuple" => {
             body.value = Expression::Tuple(Box::new(Tuple {
-                elements: vec![],
+                elements: Default::default(),
                 lpar: vec![Default::default()],
                 rpar: vec![Default::default()],
             }));
         }
         "list" => {
             body.value = Expression::List(Box::new(List {
-                elements: vec![],
+                elements: Default::default(),
                 lbracket: Default::default(),
                 rbracket: Default::default(),
-                lpar: vec![],
+                lpar: Default::default(),
-                rpar: vec![],
+                rpar: Default::default(),
             }));
         }
         "dict" => {
             body.value = Expression::Dict(Box::new(Dict {
-                elements: vec![],
+                elements: Default::default(),
                 lbrace: Default::default(),
                 rbrace: Default::default(),
-                lpar: vec![],
+                lpar: Default::default(),
-                rpar: vec![],
+                rpar: Default::default(),
             }));
         }
         _ => {
@@ -464,38 +340,13 @@ pub fn fix_unnecessary_collection_call(locator: &SourceCodeLocator, expr: &Expr)
 /// (C409) Convert `tuple([1, 2])` to `tuple(1, 2)`
 pub fn fix_unnecessary_literal_within_tuple_call(
     locator: &SourceCodeLocator,
-    expr: &Expr,
+    expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    let mut tree = match libcst_native::parse_module(
+    let mut tree = match_tree(locator, expr)?;
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    let mut body = match_expr(&mut tree)?;
-        None,
+    let call = match_call(body)?;
-    ) {
+    let arg = match_arg(call)?;
-        Ok(m) => m,
+    let (elements, whitespace_after, whitespace_before) = match &arg.value {
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
-    };
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
-    let arg = if let Some(Arg { value, .. }) = call.args.first() {
-        value
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
-    };
-    let (elements, whitespace_after, whitespace_before) = match arg {
         Expression::Tuple(inner) => (
             &inner.elements,
             &inner
@@ -544,38 +395,13 @@ pub fn fix_unnecessary_literal_within_tuple_call(
 /// (C410) Convert `list([1, 2])` to `[1, 2]`
 pub fn fix_unnecessary_literal_within_list_call(
     locator: &SourceCodeLocator,
-    expr: &Expr,
+    expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    let mut tree = match libcst_native::parse_module(
+    let mut tree = match_tree(locator, expr)?;
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    let mut body = match_expr(&mut tree)?;
-        None,
+    let call = match_call(body)?;
-    ) {
+    let arg = match_arg(call)?;
-        Ok(m) => m,
+    let (elements, whitespace_after, whitespace_before) = match &arg.value {
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
-    };
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
-    let arg = if let Some(Arg { value, .. }) = call.args.first() {
-        value
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
-    };
-    let (elements, whitespace_after, whitespace_before) = match arg {
         Expression::Tuple(inner) => (
             &inner.elements,
             &inner
@@ -624,40 +450,17 @@ pub fn fix_unnecessary_literal_within_list_call(
 }
 
 /// (C411) Convert `list([i for i in x])` to `[i for i in x]`.
-pub fn fix_unnecessary_list_call(locator: &SourceCodeLocator, expr: &Expr) -> Result<Fix> {
+pub fn fix_unnecessary_list_call(
-    // Module(SimpleStatementLine(Expr(Call(List|Tuple)))) ->
+    locator: &SourceCodeLocator,
-    // Module(SimpleStatementLine(Expr(List|Tuple)))
+    expr: &rustpython_ast::Expr,
-    let mut tree = match libcst_native::parse_module(
+) -> Result<Fix> {
-        locator.slice_source_code_range(&Range::from_located(expr)),
+    // Expr(Call(List|Tuple)))) -> Expr(List|Tuple)))
-        None,
+    let mut tree = match_tree(locator, expr)?;
-    ) {
+    let mut body = match_expr(&mut tree)?;
-        Ok(m) => m,
+    let call = match_call(body)?;
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
+    let arg = match_arg(call)?;
-    };
-    let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Statement::Simple."));
-    };
-    let body = if let Some(SmallStatement::Expr(body)) = body.body.first_mut() {
-        body
-    } else {
-        return Err(anyhow::anyhow!(
-            "Expected node to be: SmallStatement::Expr."
-        ));
-    };
-    let call = if let Expression::Call(call) = &body.value {
-        call
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Expression::Call."));
-    };
-    let arg = if let Some(Arg { value, .. }) = call.args.first() {
-        value
-    } else {
-        return Err(anyhow::anyhow!("Expected node to be: Arg."));
-    };
 
-    body.value = arg.clone();
+    body.value = arg.value.clone();
 
     let mut state = Default::default();
     tree.codegen(&mut state);

src/pyflakes/fixes.rs

@@ -2,9 +2,9 @@ use anyhow::Result;
 use libcst_native::{Codegen, ImportNames, NameOrAttribute, SmallStatement, Statement};
 use rustpython_ast::Stmt;
 
-use crate::ast::types::Range;
 use crate::autofix::{helpers, Fix};
 use crate::cst::helpers::compose_module_path;
+use crate::cst::matchers::match_tree;
 use crate::source_code_locator::SourceCodeLocator;
 
 /// Generate a Fix to remove any unused imports from an `import` statement.
@@ -15,13 +15,7 @@ pub fn remove_unused_imports(
     parent: Option<&Stmt>,
     deleted: &[&Stmt],
 ) -> Result<Fix> {
-    let mut tree = match libcst_native::parse_module(
+    let mut tree = match_tree(locator, stmt)?;
-        locator.slice_source_code_range(&Range::from_located(stmt)),
-        None,
-    ) {
-        Ok(m) => m,
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
-    };
 
     let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
         body
@@ -78,13 +72,7 @@ pub fn remove_unused_import_froms(
     parent: Option<&Stmt>,
     deleted: &[&Stmt],
 ) -> Result<Fix> {
-    let mut tree = match libcst_native::parse_module(
+    let mut tree = match_tree(locator, stmt)?;
-        locator.slice_source_code_range(&Range::from_located(stmt)),
-        None,
-    ) {
-        Ok(m) => m,
-        Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
-    };
 
     let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() {
         body