From 702690cde9d8ea701fd032df407c210550b6a7b5 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sat, 18 May 2024 23:12:52 -0400
Subject: [PATCH] Respect operator precedence in FURB110

---
 .../resources/test/fixtures/refurb/FURB110.py |  9 ++++
 .../rules/if_exp_instead_of_or_operator.rs    | 52 ++++++++++++-------
 ...es__refurb__tests__FURB110_FURB110.py.snap | 30 +++++++++++
 3 files changed, 73 insertions(+), 18 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/refurb/FURB110.py b/crates/ruff_linter/resources/test/fixtures/refurb/FURB110.py
index 706d78a4f0103..6efd0aaee4d3c 100644
--- a/crates/ruff_linter/resources/test/fixtures/refurb/FURB110.py
+++ b/crates/ruff_linter/resources/test/fixtures/refurb/FURB110.py
@@ -38,3 +38,12 @@
     else
     y
 )
+
+# FURB110
+z = (
+    x
+    if x
+    else y
+    if y > 0
+    else None
+)
diff --git a/crates/ruff_linter/src/rules/refurb/rules/if_exp_instead_of_or_operator.rs b/crates/ruff_linter/src/rules/refurb/rules/if_exp_instead_of_or_operator.rs
index 5fc89c6a149ab..131b8eb789f36 100644
--- a/crates/ruff_linter/src/rules/refurb/rules/if_exp_instead_of_or_operator.rs
+++ b/crates/ruff_linter/src/rules/refurb/rules/if_exp_instead_of_or_operator.rs
@@ -1,9 +1,14 @@
+use std::borrow::Cow;
+
 use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast as ast;
 use ruff_python_ast::comparable::ComparableExpr;
 use ruff_python_ast::helpers::contains_effect;
 use ruff_python_ast::parenthesize::parenthesized_range;
+use ruff_python_ast::Expr;
+use ruff_python_index::Indexer;
+use ruff_source_file::Locator;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -64,29 +69,13 @@ pub(crate) fn if_exp_instead_of_or_operator(checker: &mut Checker, if_expr: &ast
 
     let mut diagnostic = Diagnostic::new(IfExpInsteadOfOrOperator, *range);
 
-    // Grab the range of the `test` and `orelse` expressions.
-    let left = parenthesized_range(
-        test.into(),
-        if_expr.into(),
-        checker.indexer().comment_ranges(),
-        checker.locator().contents(),
-    )
-    .unwrap_or(test.range());
-    let right = parenthesized_range(
-        orelse.into(),
-        if_expr.into(),
-        checker.indexer().comment_ranges(),
-        checker.locator().contents(),
-    )
-    .unwrap_or(orelse.range());
-
     // Replace with `{test} or {orelse}`.
     diagnostic.set_fix(Fix::applicable_edit(
         Edit::range_replacement(
             format!(
                 "{} or {}",
-                checker.locator().slice(left),
-                checker.locator().slice(right),
+                parenthesize_test(test, if_expr, checker.indexer(), checker.locator()),
+                parenthesize_test(orelse, if_expr, checker.indexer(), checker.locator()),
             ),
             if_expr.range(),
         ),
@@ -99,3 +88,30 @@ pub(crate) fn if_exp_instead_of_or_operator(checker: &mut Checker, if_expr: &ast
 
     checker.diagnostics.push(diagnostic);
 }
+
+/// Parenthesize an expression for use in an `or` operator (e.g., parenthesize `x` in `x or y`),
+/// if it's required to maintain the correct order of operations.
+///
+/// If the expression is already parenthesized, it will be returned as-is regardless of whether
+/// the parentheses are required.
+///
+/// See: <https://docs.python.org/3/reference/expressions.html#operator-precedence>
+fn parenthesize_test<'a>(
+    expr: &Expr,
+    if_expr: &ast::ExprIf,
+    indexer: &Indexer,
+    locator: &Locator<'a>,
+) -> Cow<'a, str> {
+    if let Some(range) = parenthesized_range(
+        expr.into(),
+        if_expr.into(),
+        indexer.comment_ranges(),
+        locator.contents(),
+    ) {
+        Cow::Borrowed(locator.slice(range))
+    } else if matches!(expr, Expr::If(_) | Expr::Lambda(_) | Expr::Named(_)) {
+        Cow::Owned(format!("({})", locator.slice(expr.range())))
+    } else {
+        Cow::Borrowed(locator.slice(expr.range()))
+    }
+}
diff --git a/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB110_FURB110.py.snap b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB110_FURB110.py.snap
index 69ae14f9401c5..29e977193e329 100644
--- a/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB110_FURB110.py.snap
+++ b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB110_FURB110.py.snap
@@ -177,3 +177,33 @@ FURB110.py:34:5: FURB110 [*] Replace ternary `if` expression with `or` operator
 39    |-    y
    34 |+    x or y
 40 35 | )
+41 36 | 
+42 37 | # FURB110
+
+FURB110.py:44:5: FURB110 [*] Replace ternary `if` expression with `or` operator
+   |
+42 |   # FURB110
+43 |   z = (
+44 |       x
+   |  _____^
+45 | |     if x
+46 | |     else y
+47 | |     if y > 0
+48 | |     else None
+   | |_____________^ FURB110
+49 |   )
+   |
+   = help: Replace with `or` operator
+
+ℹ Safe fix
+41 41 | 
+42 42 | # FURB110
+43 43 | z = (
+44    |-    x
+45    |-    if x
+46    |-    else y
+   44 |+    x or (y
+47 45 |     if y > 0
+48    |-    else None
+   46 |+    else None)
+49 47 | )