Avoid assert() to assert statement conversion in expressions (#3062)

astral-sh · Feb 20, 2023 · 41f163f · 41f163f
1 parent d21dd99
commit 41f163f
Show file tree

Hide file tree

Showing 3 changed files with 163 additions and 83 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_pytest_style/PT009.py b/crates/ruff/resources/test/fixtures/flake8_pytest_style/PT009.py
@@ -17,6 +17,12 @@ def test_assert_true(self):
         self.assertTrue(**{"expr": expr, "msg": msg})  # Error, unfixable
         self.assertTrue(msg=msg, expr=expr, unexpected_arg=False)  # Error, unfixable
         self.assertTrue(msg=msg)  # Error, unfixable
+        (
+            self.assertIsNotNone(value)  # Error, unfixable
+            if expect_condition
+            else self.assertIsNone(value)  # Error, unfixable
+        )
+        return self.assertEqual(True,  False)  # Error, unfixable
 
     def test_assert_false(self):
         self.assertFalse(True)  # Error

diff --git a/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs b/crates/ruff/src/rules/flake8_pytest_style/rules/assertion.rs
@@ -12,7 +12,7 @@ use crate::checkers::ast::Checker;
 use crate::fix::Fix;
 use crate::registry::Diagnostic;
 use crate::source_code::Stylist;
-use crate::violation::{AlwaysAutofixableViolation, AutofixKind, Availability, Violation};
+use crate::violation::{AutofixKind, Availability, Violation};
 
 use super::helpers::is_falsy_constant;
 use super::unittest_assert::UnittestAssert;
@@ -94,18 +94,23 @@ impl Violation for AssertAlwaysFalse {
 define_violation!(
     pub struct UnittestAssertion {
         pub assertion: String,
+        pub fixable: bool,
     }
 );
-impl AlwaysAutofixableViolation for UnittestAssertion {
+impl Violation for UnittestAssertion {
+    const AUTOFIX: Option<AutofixKind> = Some(AutofixKind::new(Availability::Sometimes));
+
     #[derive_message_formats]
     fn message(&self) -> String {
-        let UnittestAssertion { assertion } = self;
+        let UnittestAssertion { assertion, .. } = self;
         format!("Use a regular `assert` instead of unittest-style `{assertion}`")
     }
 
-    fn autofix_title(&self) -> String {
-        let UnittestAssertion { assertion } = self;
-        format!("Replace `{assertion}(...)` with `assert ...`")
+    fn autofix_title_formatter(&self) -> Option<fn(&Self) -> String> {
+        self.fixable
+            .then_some(|UnittestAssertion { assertion, .. }| {
+                format!("Replace `{assertion}(...)` with `assert ...`")
+            })
     }
 }
 
@@ -181,13 +186,18 @@ pub fn unittest_assertion(
     match &func.node {
         ExprKind::Attribute { attr, .. } => {
             if let Ok(unittest_assert) = UnittestAssert::try_from(attr.as_str()) {
+                // We're converting an expression to a statement, so avoid applying the fix if
+                // the assertion is part of a larger expression.
+                let fixable = checker.current_expr_parent().is_none()
+                    && matches!(checker.current_stmt().node, StmtKind::Expr { .. });
                 let mut diagnostic = Diagnostic::new(
                     UnittestAssertion {
                         assertion: unittest_assert.to_string(),
+                        fixable,
                     },
                     Range::from_located(func),
                 );
-                if checker.patch(diagnostic.kind.rule()) {
+                if fixable && checker.patch(diagnostic.kind.rule()) {
                     if let Ok(stmt) = unittest_assert.generate_assert(args, keywords) {
                         diagnostic.amend(Fix::replacement(
                             unparse_stmt(&stmt, checker.stylist),