ruff
94e61100 - [`airflow`] Implement `task-branch-as-short-circuit` (`AIR004`) (#23579)

Commit
171 days ago
[`airflow`] Implement `task-branch-as-short-circuit` (`AIR004`) (#23579) ## Summary Adds a new rule `AIR004` that detects `@task.branch` decorated functions that could be replaced with `@task.short_circuit`. In Airflow, [`@task.branch`](https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/dags.html#branching) selects which downstream tasks to run by returning a list of task IDs (or an empty list to skip all). When the function has at least two `return` statements and exactly one of them returns a non-empty list, it is effectively acting as a boolean short-circuit (i.e. either run one specific set of downstream tasks or skip them all). In that case, [`@task.short_circuit`](https://www.astronomer.io/docs/learn/airflow-branch-operator#taskshort_circuit-shortcircuitoperator) is a simpler and more readable alternative that returns `True`/`False` instead. ```python # Before (AIR004) @task.branch def my_task(): if condition: return ["my_downstream_task"] return [] # After @task.short_circuit def my_task(): return condition ``` ### Implementation details - Resolves the `@task.branch` decorator via the semantic model (`airflow.decorators.task` + `.branch` attribute), handling both `@task.branch` and `@task.branch()` call forms via `map_callable`. - Uses `ReturnStatementVisitor` to collect all `return` statements recursively (including those inside nested `if`/`else`/`for`/`while` blocks). - Flags the function when: `len(returns) >= 2` and exactly one return has a non-empty list value. ### What it does NOT flag - Functions with multiple non-empty list returns (genuine branching logic). - Functions with all-empty returns (no downstream tasks selected at all). - Functions with only a single return statement. - Functions not decorated with `@task.branch`. - Functions returning non-list values (strings, `None`, etc.). ## Test Plan <!-- How was it tested? --> Added snapshot tests in `AIR004.py` covering both violation and non-violation cases: - two returns with one non-empty list - three returns with one non-empty list - nested returns - multiple non-empty returns - all-empty returns - single return - undecorated functions - `@task.short_circuit` decorated functions
Author
Parents
Loading