Do not Change Quotation Style for PT006 Autofix (#1600)

resources/test/fixtures/flake8_pytest_style/PT006.py

@@ -16,6 +16,11 @@ def test_csv_with_whitespace(param1, param2):
     ...
 
 
+@pytest.mark.parametrize("param1,param2", [(1, 2), (3, 4)])
+def test_csv_bad_quotes(param1, param2):
+    ...
+
+
 @pytest.mark.parametrize(("param1", "param2"), [(1, 2), (3, 4)])
 def test_tuple(param1, param2):
     ...

src/flake8_pytest_style/plugins/parametrize.rs

@@ -1,11 +1,14 @@
-use rustpython_ast::{Constant, Expr, ExprKind};
+use log::error;
+use rustpython_ast::{Constant, Expr, ExprContext, ExprKind};
 
 use super::helpers::is_pytest_parametrize;
+use crate::ast::helpers::create_expr;
 use crate::ast::types::Range;
 use crate::autofix::Fix;
 use crate::checkers::ast::Checker;
 use crate::flake8_pytest_style::types;
 use crate::registry::{Check, CheckCode, CheckKind};
+use crate::source_code_generator::SourceCodeGenerator;
 
 fn get_parametrize_decorator<'a>(checker: &Checker, decorators: &'a [Expr]) -> Option<&'a Expr> {
     decorators
@@ -44,11 +47,39 @@ fn check_names(checker: &mut Checker, expr: &Expr) {
                             Range::from_located(expr),
                         );
                         if checker.patch(check.kind.code()) {
-                            check.amend(Fix::replacement(
+                            let mut generator = SourceCodeGenerator::new(
-                                strings_to_python_tuple(&names),
+                                checker.style.indentation(),
-                                expr.location,
+                                checker.style.quote(),
-                                expr.end_location.unwrap(),
+                                checker.style.line_ending(),
-                            ));
+                            );
+                            generator.unparse_expr(
+                                &create_expr(ExprKind::Tuple {
+                                    elts: names
+                                        .iter()
+                                        .map(|&name| {
+                                            create_expr(ExprKind::Constant {
+                                                value: Constant::Str(name.to_string()),
+                                                kind: None,
+                                            })
+                                        })
+                                        .collect(),
+                                    ctx: ExprContext::Load,
+                                }),
+                                1,
+                            );
+                            match generator.generate() {
+                                Ok(content) => {
+                                    check.amend(Fix::replacement(
+                                        content,
+                                        expr.location,
+                                        expr.end_location.unwrap(),
+                                    ));
+                                }
+                                Err(e) => error!(
+                                    "Failed to fix wrong name(s) type in \
+                                     `@pytest.mark.parametrize`: {e}"
+                                ),
+                            };
                         }
                         checker.add_check(check);
                     }
@@ -58,11 +89,39 @@ fn check_names(checker: &mut Checker, expr: &Expr) {
                             Range::from_located(expr),
                         );
                         if checker.patch(check.kind.code()) {
-                            check.amend(Fix::replacement(
+                            let mut generator = SourceCodeGenerator::new(
-                                strings_to_python_list(&names),
+                                checker.style.indentation(),
-                                expr.location,
+                                checker.style.quote(),
-                                expr.end_location.unwrap(),
+                                checker.style.line_ending(),
-                            ));
+                            );
+                            generator.unparse_expr(
+                                &create_expr(ExprKind::List {
+                                    elts: names
+                                        .iter()
+                                        .map(|&name| {
+                                            create_expr(ExprKind::Constant {
+                                                value: Constant::Str(name.to_string()),
+                                                kind: None,
+                                            })
+                                        })
+                                        .collect(),
+                                    ctx: ExprContext::Load,
+                                }),
+                                0,
+                            );
+                            match generator.generate() {
+                                Ok(content) => {
+                                    check.amend(Fix::replacement(
+                                        content,
+                                        expr.location,
+                                        expr.end_location.unwrap(),
+                                    ));
+                                }
+                                Err(e) => error!(
+                                    "Failed to fix wrong name(s) type in \
+                                     `@pytest.mark.parametrize`: {e}"
+                                ),
+                            };
                         }
                         checker.add_check(check);
                     }
@@ -135,42 +194,28 @@ fn handle_single_name(checker: &mut Checker, expr: &Expr, value: &Expr) {
         CheckKind::ParametrizeNamesWrongType(types::ParametrizeNameType::CSV),
         Range::from_located(expr),
     );
-    if let ExprKind::Constant {
+
-        value: Constant::Str(string),
+    if checker.patch(check.kind.code()) {
-        ..
+        let mut generator = SourceCodeGenerator::new(
-    } = &value.node
+            checker.style.indentation(),
-    {
+            checker.style.quote(),
-        if checker.patch(check.kind.code()) {
+            checker.style.line_ending(),
-            check.amend(Fix::replacement(
+        );
-                format!("\"{string}\""),
+        generator.unparse_expr(&create_expr(value.node.clone()), 0);
-                expr.location,
+        match generator.generate() {
-                expr.end_location.unwrap(),
+            Ok(content) => {
-            ));
+                check.amend(Fix::replacement(
-        }
+                    content,
+                    expr.location,
+                    expr.end_location.unwrap(),
+                ));
+            }
+            Err(e) => error!("Failed to fix wrong name(s) type in `@pytest.mark.parametrize`: {e}"),
+        };
     }
     checker.add_check(check);
 }
 
-fn strings_to_python_tuple(strings: &[&str]) -> String {
-    let result = strings
-        .iter()
-        .map(|s| format!("\"{s}\""))
-        .collect::<Vec<String>>()
-        .join(", ");
-
-    format!("({result})")
-}
-
-fn strings_to_python_list(strings: &[&str]) -> String {
-    let result = strings
-        .iter()
-        .map(|s| format!("\"{s}\""))
-        .collect::<Vec<String>>()
-        .join(", ");
-
-    format!("[{result}]")
-}
-
 fn handle_value_rows(
     checker: &mut Checker,
     elts: &[Expr],

src/flake8_pytest_style/snapshots/ruff__flake8_pytest_style__tests__PT006_csv.snap

@@ -5,55 +5,55 @@ expression: checks
 - kind:
     ParametrizeNamesWrongType: csv
   location:
-    row: 19
+    row: 24
     column: 25
   end_location:
-    row: 19
+    row: 24
     column: 45
   fix: ~
   parent: ~
 - kind:
     ParametrizeNamesWrongType: csv
   location:
-    row: 24
+    row: 29
     column: 25
   end_location:
-    row: 24
+    row: 29
     column: 36
   fix:
     content: "\"param1\""
     location:
-      row: 24
+      row: 29
       column: 25
     end_location:
-      row: 24
+      row: 29
       column: 36
   parent: ~
 - kind:
     ParametrizeNamesWrongType: csv
   location:
-    row: 29
+    row: 34
     column: 25
   end_location:
-    row: 29
+    row: 34
     column: 45
   fix: ~
   parent: ~
 - kind:
     ParametrizeNamesWrongType: csv
   location:
-    row: 34
+    row: 39
     column: 25
   end_location:
-    row: 34
+    row: 39
     column: 35
   fix:
     content: "\"param1\""
     location:
-      row: 34
+      row: 39
       column: 25
     end_location:
-      row: 34
+      row: 39
       column: 35
   parent: ~
 

src/flake8_pytest_style/snapshots/ruff__flake8_pytest_style__tests__PT006_default.snap

@@ -37,47 +37,64 @@ expression: checks
       column: 56
   parent: ~
 - kind:
-    ParametrizeNamesWrongType: csv
+    ParametrizeNamesWrongType: tuple
   location:
-    row: 24
+    row: 19
     column: 25
   end_location:
-    row: 24
+    row: 19
+    column: 40
+  fix:
+    content: "(\"param1\", \"param2\")"
+    location:
+      row: 19
+      column: 25
+    end_location:
+      row: 19
+      column: 40
+  parent: ~
+- kind:
+    ParametrizeNamesWrongType: csv
+  location:
+    row: 29
+    column: 25
+  end_location:
+    row: 29
     column: 36
   fix:
     content: "\"param1\""
     location:
-      row: 24
+      row: 29
       column: 25
     end_location:
-      row: 24
+      row: 29
       column: 36
   parent: ~
 - kind:
     ParametrizeNamesWrongType: tuple
   location:
-    row: 29
+    row: 34
     column: 25
   end_location:
-    row: 29
+    row: 34
     column: 45
   fix: ~
   parent: ~
 - kind:
     ParametrizeNamesWrongType: csv
   location:
-    row: 34
+    row: 39
     column: 25
   end_location:
-    row: 34
+    row: 39
     column: 35
   fix:
     content: "\"param1\""
     location:
-      row: 34
+      row: 39
       column: 25
     end_location:
-      row: 34
+      row: 39
       column: 35
   parent: ~
 

src/flake8_pytest_style/snapshots/ruff__flake8_pytest_style__tests__PT006_list.snap

@@ -43,41 +43,58 @@ expression: checks
     column: 25
   end_location:
     row: 19
+    column: 40
+  fix:
+    content: "[\"param1\", \"param2\"]"
+    location:
+      row: 19
+      column: 25
+    end_location:
+      row: 19
+      column: 40
+  parent: ~
+- kind:
+    ParametrizeNamesWrongType: list
+  location:
+    row: 24
+    column: 25
+  end_location:
+    row: 24
     column: 45
   fix: ~
   parent: ~
 - kind:
     ParametrizeNamesWrongType: csv
   location:
-    row: 24
+    row: 29
     column: 25
   end_location:
-    row: 24
+    row: 29
     column: 36
   fix:
     content: "\"param1\""
     location:
-      row: 24
+      row: 29
       column: 25
     end_location:
-      row: 24
+      row: 29
       column: 36
   parent: ~
 - kind:
     ParametrizeNamesWrongType: csv
   location:
-    row: 34
+    row: 39
     column: 25
   end_location:
-    row: 34
+    row: 39
     column: 35
   fix:
     content: "\"param1\""
     location:
-      row: 34
+      row: 39
       column: 25
     end_location:
-      row: 34
+      row: 39
       column: 35
   parent: ~