Use Walrus Operator in if statementHas Fix

• Playground Link

Description

The walrus operator (:=) introduced in Python 3.8 allows you to assign values to variables as part of an expression. This rule aims to simplify code by using the walrus operator in if statements.

This first part of the rule identifies cases where a variable is assigned a value and then immediately used in an if statement to control flow.

id: use-walrus-operator
language: python
rule:
  pattern: "if $VAR: $$$B"
  follows:
    pattern:
      context: $VAR = $$$EXPR
      selector: expression_statement
fix: |-
  if $VAR := $$$EXPR:
    $$$B

The pattern clause finds an if statement that checks the truthiness of $VAR. If this pattern follows an expression statement where $VAR is assigned $$$EXPR, the fix clause changes the if statements to use the walrus operator.

The second part of the rule:

id: remove-declaration
rule:
  pattern:
    context: $VAR = $$$EXPR
    selector: expression_statement
  precedes:
    pattern: "if $VAR: $$$B"
fix: ''

This rule removes the standalone variable assignment when it directly precedes an if statement that uses the walrus operator. Since the assignment is now part of the if statement, the separate declaration is no longer needed.

By applying these rules, you can refactor your Python code to be more concise and readable, taking advantage of the walrus operator's ability to combine an assignment with an expression.

YAML

id: use-walrus-operator
language: python
rule:
  follows:
    pattern:
      context: $VAR = $$$EXPR
      selector: expression_statement
  pattern: "if $VAR: $$$B"
fix: |-
  if $VAR := $$$EXPR:
    $$$B
---
id: remove-declaration
language: python
rule:
  pattern:
    context: $VAR = $$$EXPR
    selector: expression_statement
  precedes:
    pattern: "if $VAR: $$$B"
fix: ''

Example

a = foo()

if a:
    do_bar()

Diff

a = foo() 

if a: 
if a := foo(): 
    do_bar()

Contributed by

Inspired by reddit user /u/jackerhack