[ruff] Avoid emitting assignment-in-assert when all references to the assigned variable are themselves inside asserts (RUF018)

crates/ruff_linter/resources/test/fixtures/ruff/RUF018.py

@@ -1,7 +1,20 @@
 # RUF018
 assert (x := 0) == 0
 assert x, (y := "error")
+print(x, y)
 
 # OK
 if z := 0:
     pass
+
+
+# These should not be flagged, because the only uses of the variables defined
+# are themselves within `assert` statements.
+
+# Here the `t` variable is referenced
+# from a later `assert` statement:
+assert (t:=cancel((F, G))) == (1, P, Q)
+assert isinstance(t, tuple)
+
+# Here the `g` variable is referenced from within the same `assert` statement:
+assert (g:=solve(groebner(eqs, s), dict=True)) == sol, g

crates/ruff_linter/src/checkers/ast/analyze/bindings.rs

@@ -18,6 +18,7 @@ pub(crate) fn bindings(checker: &mut Checker) {
         Rule::UnsortedDunderSlots,
         Rule::UnusedVariable,
         Rule::UnquotedTypeAlias,
+        Rule::AssignmentInAssert,
     ]) {
         return;
     }
@@ -87,5 +88,10 @@ pub(crate) fn bindings(checker: &mut Checker) {
                 checker.diagnostics.push(diagnostic);
             }
         }
+        if checker.enabled(Rule::AssignmentInAssert) {
+            if let Some(diagnostic) = ruff::rules::assignment_in_assert(checker, binding) {
+                checker.diagnostics.push(diagnostic);
+            }
+        }
     }
 }

crates/ruff_linter/src/checkers/ast/analyze/expression.rs

@@ -1648,11 +1648,6 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 ruff::rules::parenthesize_chained_logical_operators(checker, bool_op);
             }
         }
-        Expr::Named(..) => {
-            if checker.enabled(Rule::AssignmentInAssert) {
-                ruff::rules::assignment_in_assert(checker, expr);
-            }
-        }
         Expr::Lambda(lambda_expr) => {
             if checker.enabled(Rule::ReimplementedOperator) {
                 refurb::rules::reimplemented_operator(checker, &lambda_expr.into());

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -971,10 +971,13 @@ impl<'a> Visitor<'a> for Checker<'a> {
                 msg,
                 range: _,
             }) => {
+                let snapshot = self.semantic.flags;
+                self.semantic.flags |= SemanticModelFlags::ASSERT_STATEMENT;
                 self.visit_boolean_test(test);
                 if let Some(expr) = msg {
                     self.visit_expr(expr);
                 }
+                self.semantic.flags = snapshot;
             }
             Stmt::With(ast::StmtWith {
                 items,
@@ -1954,6 +1957,9 @@ impl<'a> Checker<'a> {
         if self.semantic.in_exception_handler() {
             flags |= BindingFlags::IN_EXCEPT_HANDLER;
         }
+        if self.semantic.in_assert_statement() {
+            flags |= BindingFlags::IN_ASSERT_STATEMENT;
+        }
 
         // Create the `Binding`.
         let binding_id = self.semantic.push_binding(range, kind, flags);

crates/ruff_linter/src/rules/ruff/rules/assignment_in_assert.rs

@@ -1,7 +1,6 @@
-use ruff_python_ast::Expr;
-
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
+use ruff_python_semantic::Binding;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -23,12 +22,27 @@ use crate::checkers::ast::Checker;
 /// ## Examples
 /// ```python
 /// assert (x := 0) == 0
+/// print(x)
 /// ```
 ///
 /// Use instead:
 /// ```python
 /// x = 0
 /// assert x == 0
+/// print(x)
+/// ```
+///
+/// The rule avoids flagging named expressions that define variables which are
+/// only referenced from inside `assert` statements; the following will not
+/// trigger the rule:
+/// ```python
+/// assert (x := y**2) > 42, f"Expected >42 but got {x}"
+/// ```
+///
+/// Nor will this:
+/// ```python
+/// assert (x := y**2) > 42
+/// assert x < 1_000_000
 /// ```
 ///
 /// ## References
@@ -44,10 +58,24 @@ impl Violation for AssignmentInAssert {
 }
 
 /// RUF018
-pub(crate) fn assignment_in_assert(checker: &mut Checker, value: &Expr) {
-    if checker.semantic().current_statement().is_assert_stmt() {
-        checker
-            .diagnostics
-            .push(Diagnostic::new(AssignmentInAssert, value.range()));
+pub(crate) fn assignment_in_assert(checker: &Checker, binding: &Binding) -> Option<Diagnostic> {
+    if !binding.in_assert_statement() {
+        return None;
+    }
+
+    let semantic = checker.semantic();
+
+    let parent_expression = binding.expression(semantic)?.as_named_expr()?;
+
+    if binding
+        .references()
+        .all(|reference| semantic.reference(reference).in_assert_statement())
+    {
+        return None;
     }
+
+    Some(Diagnostic::new(
+        AssignmentInAssert,
+        parent_expression.range(),
+    ))
 }

crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF018_RUF018.py.snap

@@ -1,13 +1,13 @@
 ---
 source: crates/ruff_linter/src/rules/ruff/mod.rs
-snapshot_kind: text
 ---
 RUF018.py:2:9: RUF018 Avoid assignment expressions in `assert` statements
   |
 1 | # RUF018
 2 | assert (x := 0) == 0
   |         ^^^^^^ RUF018
 3 | assert x, (y := "error")
+4 | print(x, y)
   |
 
 RUF018.py:3:12: RUF018 Avoid assignment expressions in `assert` statements
@@ -16,6 +16,5 @@ RUF018.py:3:12: RUF018 Avoid assignment expressions in `assert` statements
 2 | assert (x := 0) == 0
 3 | assert x, (y := "error")
   |            ^^^^^^^^^^^^ RUF018
-4 | 
-5 | # OK
+4 | print(x, y)
   |

crates/ruff_python_semantic/src/binding.rs

@@ -135,6 +135,15 @@ impl<'a> Binding<'a> {
         self.flags.contains(BindingFlags::IN_EXCEPT_HANDLER)
     }
 
+    /// Return `true` if this [`Binding`] took place inside an `assert` statement,
+    /// e.g. `y` in:
+    /// ```python
+    /// assert (y := x**2), y
+    /// ```
+    pub const fn in_assert_statement(&self) -> bool {
+        self.flags.contains(BindingFlags::IN_ASSERT_STATEMENT)
+    }
+
     /// Return `true` if this [`Binding`] represents a [PEP 613] type alias
     /// e.g. `OptString` in:
     /// ```python
@@ -266,6 +275,15 @@ impl<'a> Binding<'a> {
             .map(|statement_id| semantic.statement(statement_id))
     }
 
+    /// Returns the expression in which the binding was defined
+    /// (e.g. for the binding `x` in `y = (x := 1)`, return the node representing `x := 1`).
+    ///
+    /// This is only really applicable for assignment expressions.
+    pub fn expression<'b>(&self, semantic: &SemanticModel<'b>) -> Option<&'b ast::Expr> {
+        self.source
+            .and_then(|expression_id| semantic.parent_expression(expression_id))
+    }
+
     /// Returns the range of the binding's parent.
     pub fn parent_range(&self, semantic: &SemanticModel) -> Option<TextRange> {
         self.statement(semantic).and_then(|parent| {
@@ -406,6 +424,14 @@ bitflags! {
         /// [PEP 695]: https://peps.python.org/pep-0695/#generic-type-alias
         const DEFERRED_TYPE_ALIAS = 1 << 12;
 
+        /// The binding took place inside an `assert` statement
+        ///
+        /// For example, `x` in the following snippet:
+        /// ```python
+        /// assert (x := y**2) > 42, x
+        /// ```
+        const IN_ASSERT_STATEMENT = 1 << 13;
+
         /// The binding represents any type alias.
         const TYPE_ALIAS = Self::ANNOTATED_TYPE_ALIAS.bits() | Self::DEFERRED_TYPE_ALIAS.bits();
     }

crates/ruff_python_semantic/src/model.rs

@@ -1839,6 +1839,11 @@ impl<'a> SemanticModel<'a> {
         self.flags.intersects(SemanticModelFlags::EXCEPTION_HANDLER)
     }
 
+    /// Return `true` if the model is in an `assert` statement.
+    pub const fn in_assert_statement(&self) -> bool {
+        self.flags.intersects(SemanticModelFlags::ASSERT_STATEMENT)
+    }
+
     /// Return `true` if the model is in an f-string.
     pub const fn in_f_string(&self) -> bool {
         self.flags.intersects(SemanticModelFlags::F_STRING)
@@ -2432,6 +2437,14 @@ bitflags! {
         /// [PEP 695]: https://peps.python.org/pep-0695/#generic-type-alias
         const DEFERRED_TYPE_ALIAS = 1 << 28;
 
+        /// The model is visiting an `assert` statement.
+        ///
+        /// For example, the model might be visiting `y` in
+        /// ```python
+        /// assert (y := x**2) > 42, y
+        /// ```
+        const ASSERT_STATEMENT = 1 << 29;
+
         /// The context is in any type annotation.
         const ANNOTATION = Self::TYPING_ONLY_ANNOTATION.bits() | Self::RUNTIME_EVALUATED_ANNOTATION.bits() | Self::RUNTIME_REQUIRED_ANNOTATION.bits();
 

crates/ruff_python_semantic/src/reference.rs

@@ -93,6 +93,11 @@ impl ResolvedReference {
         self.flags
             .intersects(SemanticModelFlags::ANNOTATED_TYPE_ALIAS)
     }
+
+    /// Return `true` if the context is inside an `assert` statement
+    pub const fn in_assert_statement(&self) -> bool {
+        self.flags.intersects(SemanticModelFlags::ASSERT_STATEMENT)
+    }
 }
 
 impl Ranged for ResolvedReference {