Implementation:Mlflow Mlflow Walrus Operator Rule

Knowledge Sources	Mlflow_Mlflow
Domains	Static Analysis, Code Style
Last Updated	2026-02-13 20:00 GMT

Overview

Clint lint rule that detects assign-then-test patterns and suggests using the Python walrus operator (:=) for more concise, idiomatic code.

Description

This module implements a custom lint rule for the Clint linter that identifies opportunities to use Python's walrus operator (:=, introduced in Python 3.8). It flags the pattern where a variable is assigned and then immediately tested for truthiness in an if statement.

UseWalrusOperator is a Rule subclass containing the core detection logic in its static check() method. It validates that a pattern qualifies for the walrus operator transformation by enforcing all of the following conditions:

1. The previous statement is a simple ast.Assign (not augmented or annotated assignment)
2. The assignment is a single-line statement (not multi-line)
3. There is exactly one target and it is a simple Name (no tuple unpacking)
4. The if condition is a direct truthiness test of the assigned variable (not a comparison)
5. The variable is used inside the if body
6. The variable is NOT used in any elif/else branches
7. The variable is NOT used in any following statements after the if block
8. The resulting fixed line would not exceed 100 characters (calculated by measuring the original value width and the resulting if var := value: length)

WalrusOperatorVisitor is an AST NodeVisitor that traverses all statement blocks in a Python source file to find violations. It visits:

• FunctionDef and AsyncFunctionDef bodies
• If bodies and else clauses
• For, AsyncFor, and While bodies and else clauses
• With and AsyncWith bodies
• Try bodies, exception handlers, else and finally clauses
• Match case bodies

Two helper functions support the detection:

• _name_used_in_stmts() - Checks if a variable name appears (in Load context) in a list of statements
• _name_used_in_node() - Recursively checks AST nodes using pattern matching, skipping nested function/class definitions to avoid false positives from inner scopes

Usage

This rule is automatically invoked by the Clint linter when running uv run clint . on the MLflow codebase. It produces violations for code that could be simplified from a = expr; if a: ... to if a := expr: .... The rule is designed to be conservative, only flagging clear-cut cases where the transformation is safe and does not hurt readability.

Code Reference

Source Location

• Repository: Mlflow_Mlflow
• File: dev/clint/src/clint/rules/use_walrus_operator.py
• Lines: 1-172

Signature

class UseWalrusOperator(Rule):
    def _message(self) -> str: ...

    @staticmethod
    def check(
        if_node: ast.If,
        prev_stmt: ast.stmt,
        following_stmts: list[ast.stmt],
    ) -> bool: ...

class WalrusOperatorVisitor(ast.NodeVisitor):
    def __init__(self) -> None: ...
    violations: list[ast.stmt]

    def _check_stmts(self, stmts: list[ast.stmt]) -> None: ...
    def visit_FunctionDef(self, node: ast.FunctionDef) -> None: ...
    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: ...
    def visit_If(self, node: ast.If) -> None: ...
    def visit_For(self, node: ast.For) -> None: ...
    def visit_AsyncFor(self, node: ast.AsyncFor) -> None: ...
    def visit_While(self, node: ast.While) -> None: ...
    def visit_With(self, node: ast.With) -> None: ...
    def visit_AsyncWith(self, node: ast.AsyncWith) -> None: ...
    def visit_Try(self, node: ast.Try) -> None: ...
    def visit_Match(self, node: ast.Match) -> None: ...

def _name_used_in_stmts(name: str, stmts: list[ast.stmt]) -> bool: ...
def _name_used_in_node(name: str, node: ast.AST) -> bool: ...

Import

from clint.rules.use_walrus_operator import UseWalrusOperator, WalrusOperatorVisitor

I/O Contract

Inputs

Name	Type	Required	Description
if_node	ast.If	Yes	The if-statement AST node to check
prev_stmt	ast.stmt	Yes	The statement immediately preceding the if-statement
following_stmts	list[ast.stmt]	Yes	All statements that follow the if-statement in the same block

Outputs

Name	Type	Description
bool (from check)	bool	True if the pattern qualifies for walrus operator transformation
violations	list[ast.stmt]	List of assignment statements (from WalrusOperatorVisitor) that should use the walrus operator

Usage Examples

Code That Triggers the Rule

# FLAGGED: Variable assigned then immediately tested for truthiness
result = get_data()
if result:
    process(result)

# FIX: Use walrus operator
if result := get_data():
    process(result)

Code That Does NOT Trigger the Rule

# NOT flagged: Variable used after the if block
result = get_data()
if result:
    process(result)
log(result)  # Used after if block

# NOT flagged: Comparison in condition (not a truthiness test)
result = get_data()
if result == expected:
    process(result)

# NOT flagged: Variable used in else branch
result = get_data()
if result:
    process(result)
else:
    handle_missing(result)

# NOT flagged: Multi-line assignment
result = (
    get_data()
)
if result:
    process(result)

# NOT flagged: Fixed line would exceed 100 characters
very_long_variable_name = some_very_long_function_call_with_many_arguments(arg1, arg2, arg3)
if very_long_variable_name:
    process(very_long_variable_name)

Using the Visitor Programmatically

import ast
from clint.rules.use_walrus_operator import WalrusOperatorVisitor

source = """
def example():
    x = compute()
    if x:
        use(x)
"""

tree = ast.parse(source)
visitor = WalrusOperatorVisitor()
visitor.visit(tree)

for violation in visitor.violations:
    print(f"Line {violation.lineno}: Use walrus operator")

Related Pages

• Environment:Mlflow_Mlflow_Python_Runtime_Environment