11use ruff_diagnostics:: { Diagnostic , Edit , Fix , FixAvailability , Violation } ;
22use ruff_macros:: { derive_message_formats, violation} ;
3- use ruff_python_ast:: { self as ast, ElifElseClause , Expr , Stmt } ;
3+ use ruff_python_ast:: comparable:: ComparableExpr ;
4+ use ruff_python_ast:: helpers:: contains_effect;
5+ use ruff_python_ast:: { self as ast, BoolOp , ElifElseClause , Expr , Stmt } ;
46use ruff_python_semantic:: analyze:: typing:: { is_sys_version_block, is_type_checking_block} ;
57use ruff_text_size:: { Ranged , TextRange } ;
68
@@ -9,12 +11,15 @@ use crate::fix::edits::fits;
911
1012/// ## What it does
1113/// Check for `if`-`else`-blocks that can be replaced with a ternary operator.
14+ /// Moreover, in [preview], check if these ternary expressions can be
15+ /// further simplified to binary expressions.
1216///
1317/// ## Why is this bad?
1418/// `if`-`else`-blocks that assign a value to a variable in both branches can
15- /// be expressed more concisely by using a ternary operator.
19+ /// be expressed more concisely by using a ternary or binary operator.
1620///
1721/// ## Example
22+ ///
1823/// ```python
1924/// if foo:
2025/// bar = x
@@ -27,24 +32,51 @@ use crate::fix::edits::fits;
2732/// bar = x if foo else y
2833/// ```
2934///
35+ /// Or, in [preview]:
36+ ///
37+ /// ```python
38+ /// if cond:
39+ /// z = cond
40+ /// else:
41+ /// z = other_cond
42+ /// ```
43+ ///
44+ /// Use instead:
45+ ///
46+ /// ```python
47+ /// z = cond or other_cond
48+ /// ```
49+ ///
3050/// ## References
3151/// - [Python documentation: Conditional expressions](https://docs.python.org/3/reference/expressions.html#conditional-expressions)
52+ ///
53+ /// [preview]: https://docs.astral.sh/ruff/preview/
3254#[ violation]
3355pub struct IfElseBlockInsteadOfIfExp {
56+ /// The ternary or binary expression to replace the `if`-`else`-block.
3457 contents : String ,
58+ /// Whether to use a binary or ternary assignment.
59+ kind : AssignmentKind ,
3560}
3661
3762impl Violation for IfElseBlockInsteadOfIfExp {
3863 const FIX_AVAILABILITY : FixAvailability = FixAvailability :: Sometimes ;
3964
4065 #[ derive_message_formats]
4166 fn message ( & self ) -> String {
42- let IfElseBlockInsteadOfIfExp { contents } = self ;
43- format ! ( "Use ternary operator `{contents}` instead of `if`-`else`-block" )
67+ let IfElseBlockInsteadOfIfExp { contents, kind } = self ;
68+ match kind {
69+ AssignmentKind :: Ternary => {
70+ format ! ( "Use ternary operator `{contents}` instead of `if`-`else`-block" )
71+ }
72+ AssignmentKind :: Binary => {
73+ format ! ( "Use binary operator `{contents}` instead of `if`-`else`-block" )
74+ }
75+ }
4476 }
4577
4678 fn fix_title ( & self ) -> Option < String > {
47- let IfElseBlockInsteadOfIfExp { contents } = self ;
79+ let IfElseBlockInsteadOfIfExp { contents, .. } = self ;
4880 Some ( format ! ( "Replace `if`-`else`-block with `{contents}`" ) )
4981 }
5082}
@@ -121,9 +153,59 @@ pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &a
121153 return ;
122154 }
123155
124- let target_var = & body_target;
125- let ternary = ternary ( target_var, body_value, test, else_value) ;
126- let contents = checker. generator ( ) . stmt ( & ternary) ;
156+ // In most cases we should now suggest a ternary operator,
157+ // but there are three edge cases where a binary operator
158+ // is more appropriate.
159+ //
160+ // For the reader's convenience, here is how
161+ // the notation translates to the if-else block:
162+ //
163+ // ```python
164+ // if test:
165+ // target_var = body_value
166+ // else:
167+ // target_var = else_value
168+ // ```
169+ //
170+ // The match statement below implements the following
171+ // logic:
172+ // - If `test == body_value` and preview enabled, replace with `target_var = test or else_value`
173+ // - If `test == not body_value` and preview enabled, replace with `target_var = body_value and else_value`
174+ // - If `not test == body_value` and preview enabled, replace with `target_var = body_value and else_value`
175+ // - Otherwise, replace with `target_var = body_value if test else else_value`
176+ let ( contents, assignment_kind) =
177+ match ( checker. settings . preview . is_enabled ( ) , test, body_value) {
178+ ( true , test_node, body_node)
179+ if ComparableExpr :: from ( test_node) == ComparableExpr :: from ( body_node)
180+ && !contains_effect ( test_node, |id| {
181+ checker. semantic ( ) . has_builtin_binding ( id)
182+ } ) =>
183+ {
184+ let target_var = & body_target;
185+ let binary = assignment_binary_or ( target_var, body_value, else_value) ;
186+ ( checker. generator ( ) . stmt ( & binary) , AssignmentKind :: Binary )
187+ }
188+ ( true , test_node, body_node)
189+ if ( test_node. as_unary_op_expr ( ) . is_some_and ( |op_expr| {
190+ op_expr. op . is_not ( )
191+ && ComparableExpr :: from ( & op_expr. operand ) == ComparableExpr :: from ( body_node)
192+ } ) || body_node. as_unary_op_expr ( ) . is_some_and ( |op_expr| {
193+ op_expr. op . is_not ( )
194+ && ComparableExpr :: from ( & op_expr. operand ) == ComparableExpr :: from ( test_node)
195+ } ) ) && !contains_effect ( test_node, |id| {
196+ checker. semantic ( ) . has_builtin_binding ( id)
197+ } ) =>
198+ {
199+ let target_var = & body_target;
200+ let binary = assignment_binary_and ( target_var, body_value, else_value) ;
201+ ( checker. generator ( ) . stmt ( & binary) , AssignmentKind :: Binary )
202+ }
203+ _ => {
204+ let target_var = & body_target;
205+ let ternary = assignment_ternary ( target_var, body_value, test, else_value) ;
206+ ( checker. generator ( ) . stmt ( & ternary) , AssignmentKind :: Ternary )
207+ }
208+ } ;
127209
128210 // Don't flag if the resulting expression would exceed the maximum line length.
129211 if !fits (
@@ -139,6 +221,7 @@ pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &a
139221 let mut diagnostic = Diagnostic :: new (
140222 IfElseBlockInsteadOfIfExp {
141223 contents : contents. clone ( ) ,
224+ kind : assignment_kind,
142225 } ,
143226 stmt_if. range ( ) ,
144227 ) ;
@@ -154,7 +237,18 @@ pub(crate) fn if_else_block_instead_of_if_exp(checker: &mut Checker, stmt_if: &a
154237 checker. diagnostics . push ( diagnostic) ;
155238}
156239
157- fn ternary ( target_var : & Expr , body_value : & Expr , test : & Expr , orelse_value : & Expr ) -> Stmt {
240+ #[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
241+ enum AssignmentKind {
242+ Binary ,
243+ Ternary ,
244+ }
245+
246+ fn assignment_ternary (
247+ target_var : & Expr ,
248+ body_value : & Expr ,
249+ test : & Expr ,
250+ orelse_value : & Expr ,
251+ ) -> Stmt {
158252 let node = ast:: ExprIf {
159253 test : Box :: new ( test. clone ( ) ) ,
160254 body : Box :: new ( body_value. clone ( ) ) ,
@@ -168,3 +262,33 @@ fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Exp
168262 } ;
169263 node1. into ( )
170264}
265+
266+ fn assignment_binary_and ( target_var : & Expr , left_value : & Expr , right_value : & Expr ) -> Stmt {
267+ let node = ast:: ExprBoolOp {
268+ op : BoolOp :: And ,
269+ values : vec ! [ left_value. clone( ) , right_value. clone( ) ] ,
270+ range : TextRange :: default ( ) ,
271+ } ;
272+ let node1 = ast:: StmtAssign {
273+ targets : vec ! [ target_var. clone( ) ] ,
274+ value : Box :: new ( node. into ( ) ) ,
275+ range : TextRange :: default ( ) ,
276+ } ;
277+ node1. into ( )
278+ }
279+
280+ fn assignment_binary_or ( target_var : & Expr , left_value : & Expr , right_value : & Expr ) -> Stmt {
281+ ( ast:: StmtAssign {
282+ range : TextRange :: default ( ) ,
283+ targets : vec ! [ target_var. clone( ) ] ,
284+ value : Box :: new (
285+ ( ast:: ExprBoolOp {
286+ range : TextRange :: default ( ) ,
287+ op : BoolOp :: Or ,
288+ values : vec ! [ left_value. clone( ) , right_value. clone( ) ] ,
289+ } )
290+ . into ( ) ,
291+ ) ,
292+ } )
293+ . into ( )
294+ }
0 commit comments