[airflow] Implement task-branch-as-short-circuit (AIR004) (astral-sh#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: Dev-iL

crates/ruff_linter/resources/test/fixtures/airflow/AIR004.py

@@ -0,0 +1,145 @@
+from airflow.decorators import task
+from airflow.operators.python import BranchPythonOperator, ShortCircuitOperator
+from airflow.providers.standard.operators.python import (
+    BranchPythonOperator as ProviderBranchPythonOperator,
+)
+
+condition1 = True
+condition2 = True
+condition3 = True
+
+
+# Violations (should trigger AIR004):
+
+@task.branch
+def two_returns_one_non_empty():  # AIR004
+    if condition1:
+        return ["my_downstream_task"]
+    return []
+
+
+@task.branch
+def three_returns_one_non_empty():  # AIR004
+    if condition1:
+        return []
+    if condition2:
+        return ["another_downstream_task"]
+    return []
+
+
+@task.branch
+def nested_returns_one_non_empty():  # AIR004
+    if condition1:
+        if condition2:
+            return []
+        return ["downstream_task"]
+    return []
+
+
+@task.branch()
+def with_parens():  # AIR004
+    if condition1:
+        return ["downstream_task"]
+    return []
+
+
+@task.branch
+def bare_return_and_list():  # AIR004
+    if condition1:
+        return ["downstream_task"]
+    return
+
+
+@task.branch
+def none_return_and_list():  # AIR004
+    if condition1:
+        return ["downstream_task"]
+    return None
+
+
+# No violations:
+
+@task.branch
+def two_non_empty_returns():
+    if condition1:
+        return ["task_a"]
+    if condition2:
+        return ["task_b"]
+    return []
+
+
+@task.branch
+def all_empty_returns():
+    if condition1:
+        return []
+    if condition2:
+        return []
+    return []
+
+
+@task.branch
+def single_return():
+    return ["downstream_task"]
+
+
+def not_decorated():
+    if condition1:
+        return ["downstream_task"]
+    return []
+
+
+@task.short_circuit
+def already_short_circuit():
+    if condition1:
+        return True
+    return False
+
+
+# BranchPythonOperator violations (should trigger AIR004):
+
+
+def operator_short_circuit_candidate():
+    if condition1:
+        return ["downstream_task"]
+    return []
+
+
+BranchPythonOperator(task_id="task", python_callable=operator_short_circuit_candidate)  # AIR004
+
+
+def provider_short_circuit_candidate():
+    if condition1:
+        return ["downstream_task"]
+    return []
+
+
+ProviderBranchPythonOperator(  # AIR004
+    task_id="task", python_callable=provider_short_circuit_candidate
+)
+
+
+# BranchPythonOperator non-violations:
+
+
+def operator_two_non_empty():
+    if condition1:
+        return ["task_a"]
+    if condition2:
+        return ["task_b"]
+    return []
+
+
+BranchPythonOperator(task_id="task", python_callable=operator_two_non_empty)
+
+
+def operator_single_return():
+    return ["downstream_task"]
+
+
+BranchPythonOperator(task_id="task", python_callable=operator_single_return)
+
+ShortCircuitOperator(task_id="task", python_callable=operator_short_circuit_candidate)
+
+BranchPythonOperator(task_id="task", python_callable=lambda: ["downstream_task"])
+
+BranchPythonOperator(task_id="task")

crates/ruff_linter/resources/test/fixtures/airflow/AIR004_sdk.py

@@ -0,0 +1,24 @@
+from airflow.sdk import task
+
+condition1 = True
+condition2 = True
+
+
+# Violations (should trigger AIR004):
+
+@task.branch
+def sdk_two_returns_one_non_empty():  # AIR004
+    if condition1:
+        return ["my_downstream_task"]
+    return []
+
+
+# No violations:
+
+@task.branch
+def sdk_two_non_empty_returns():
+    if condition1:
+        return ["task_a"]
+    if condition2:
+        return ["task_b"]
+    return []

crates/ruff_linter/src/checkers/ast/analyze/expression.rs

@@ -1322,6 +1322,9 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
             if checker.is_rule_enabled(Rule::Airflow3DagDynamicValue) {
                 airflow::rules::airflow_3_dag_dynamic_value(checker, call);
             }
+            if checker.is_rule_enabled(Rule::AirflowTaskBranchAsShortCircuit) {
+                airflow::rules::branch_python_operator_as_short_circuit(checker, call);
+            }
             if checker.is_rule_enabled(Rule::UnnecessaryCastToInt) {
                 ruff::rules::unnecessary_cast_to_int(checker, call);
             }

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -339,6 +339,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.is_rule_enabled(Rule::PytestParameterWithDefaultArgument) {
                 flake8_pytest_style::rules::parameter_with_default_argument(checker, function_def);
             }
+            if checker.is_rule_enabled(Rule::AirflowTaskBranchAsShortCircuit) {
+                airflow::rules::task_branch_as_short_circuit(checker, function_def);
+            }
             if checker.is_rule_enabled(Rule::Airflow3Removal) {
                 airflow::rules::airflow_3_removal_function_def(checker, function_def);
             }

crates/ruff_linter/src/codes.rs

@@ -1134,6 +1134,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Airflow, "001") => rules::airflow::rules::AirflowVariableNameTaskIdMismatch,
         (Airflow, "002") => rules::airflow::rules::AirflowDagNoScheduleArgument,
         (Airflow, "003") => rules::airflow::rules::AirflowVariableGetOutsideTask,
+        (Airflow, "004") => rules::airflow::rules::AirflowTaskBranchAsShortCircuit,
         (Airflow, "201") => rules::airflow::rules::AirflowXcomPullInTemplateString,
         (Airflow, "301") => rules::airflow::rules::Airflow3Removal,
         (Airflow, "302") => rules::airflow::rules::Airflow3MovedToProvider,

crates/ruff_linter/src/rules/airflow/helpers.rs

@@ -326,3 +326,23 @@ pub(crate) fn is_airflow_task(function_def: &StmtFunctionDef, semantic: &Semanti
         false
     })
 }
+
+/// Returns `true` if the given function is decorated with a specific `@task.<variant>`
+/// form (e.g., `@task.branch` or `@task.short_circuit`).
+pub(crate) fn is_airflow_task_variant(
+    function_def: &StmtFunctionDef,
+    semantic: &SemanticModel,
+    variant: &str,
+) -> bool {
+    function_def.decorator_list.iter().any(|decorator| {
+        let expr = map_callable(&decorator.expression);
+        if let Expr::Attribute(ExprAttribute { value, attr, .. }) = expr {
+            attr.as_str() == variant
+                && semantic.resolve_qualified_name(value).is_some_and(|qn| {
+                    matches!(qn.segments(), ["airflow", "decorators" | "sdk", "task"])
+                })
+        } else {
+            false
+        }
+    })
+}

crates/ruff_linter/src/rules/airflow/mod.rs

@@ -21,6 +21,8 @@ mod tests {
         Rule::AirflowVariableGetOutsideTask,
         Path::new("AIR003_dag_decorator.py")
     )]
+    #[test_case(Rule::AirflowTaskBranchAsShortCircuit, Path::new("AIR004.py"))]
+    #[test_case(Rule::AirflowTaskBranchAsShortCircuit, Path::new("AIR004_sdk.py"))]
     #[test_case(Rule::AirflowXcomPullInTemplateString, Path::new("AIR201.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_args.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_names.py"))]

crates/ruff_linter/src/rules/airflow/rules/mod.rs

@@ -6,6 +6,7 @@ pub(crate) use removal_in_3::*;
 pub(crate) use runtime_value_in_dag_or_task::*;
 pub(crate) use suggested_to_move_to_provider_in_3::*;
 pub(crate) use suggested_to_update_3_0::*;
+pub(crate) use task_branch_as_short_circuit::*;
 pub(crate) use task_variable_name::*;
 pub(crate) use variable_get_outside_task::*;
 pub(crate) use xcom_pull_in_template_string::*;
@@ -18,6 +19,7 @@ mod removal_in_3;
 mod runtime_value_in_dag_or_task;
 mod suggested_to_move_to_provider_in_3;
 mod suggested_to_update_3_0;
+mod task_branch_as_short_circuit;
 mod task_variable_name;
 mod variable_get_outside_task;
 mod xcom_pull_in_template_string;