[`airflow`] Implement `task-branch-as-short-circuit` (`AIR004`) by Dev-iL · Pull Request #23579 · astral-sh/ruff

Dev-iL · 2026-02-26T11:14:08Z

Summary

Adds a new rule AIR004 that detects @task.branch decorated functions that could be replaced with @task.short_circuit.

In Airflow, @task.branch 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 is a simpler and more readable alternative that returns True/False instead.

# 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

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

ntBre · 2026-02-26T13:13:46Z

@Lee-W would you mind taking a look at this?

Lee-W · 2026-02-26T13:57:33Z

Will be out these days. Will take a look next week. Thanks!

sjyangkevin

Thanks for implementing it, some small feedback

sjyangkevin · 2026-02-26T13:55:58Z

Thank you for implementing it! Just want to share some feedback about the rule and the code AIR003. Do you think it is a good idea to have a more generic name to group similar rules? I feel the rule is more like “best practice” we can suggest change. So, later on we could extend on this to add more rules, instead of making this rule specific to the branch -> short circuit

Do you mean reserving a range of codes (e.g. AIR1XX for best practices), or have one giant rule that has various best practices in it? If the latter, I would be careful with increasing the scope of a rule, to provide users the ability to selectively turn them on/off.

I can see currently AIR001 and AIR002 are also for best practices. I am just trying to think how many such rules we might introduce. I think we have a trade-off here:

use one code to implement a rule cover a best practice, users have flexibility to selectively apply the rule.

use one code to group multiple rules, and introduce (suggest fixes) all best practices at once

The second might be a better approach from my perspective, as if we would like to introduce a better way to write the code, or best practices, then probably it is not harm to introduce them all 🤔 , but I am open for different view.

Do you have a specific example in mind of a rule that does multiple different things?

We could have a "simplify branching" rule with multiple possible fixes.

I can also imagine users wanting to transition to taskflow would like a rule that flags operators with an equivalent taskflow decorator - this would be multiple checks under the same code.

Is it a problem to rename rules later if we want to expand the functionality? In my personal projects that would be a standard refactor for when code evolves.

What I saw in other rule families is a format like AIR01X - for rules dealing with topic A, AIR02X for rules on topic B, etc., but individual rules deal with a very specific pattern.

I think we should differentiate different practices so that users can configurate it

astral-sh-bot · 2026-02-26T15:03:01Z

`ruff-ecosystem` results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

Miretpl · 2026-03-03T18:53:04Z

Hi @Dev-iL, probably similar to AIR005 in some ways (probably implementation detail). As we are validating the task.branch usage, I think we should also check the non-taskflow usage of the branch operator

Dev-iL · 2026-03-03T18:55:23Z

Hi @Dev-iL, probably similar to AIR005 in some ways (probably implementation detail). As we are validating the task.branch usage, I think we should also check the non-taskflow usage of the branch operator

Sounds good. Are we going to suggest to replace it with the taskflow version?

Miretpl · 2026-03-03T19:07:35Z

Sounds good. Are we going to suggest to replace it with the taskflow version?

Personally, I think it is something that the community should decide if TaskFlow is the recommended way of writing Dags in Airflow (if it hasn't decided already).

Dev-iL · 2026-03-04T11:38:29Z

Personally, I think it is something that the community should decide if TaskFlow is the recommended way of writing Dags in Airflow (if it hasn't decided already).

I think there should be a separate rule to assist with migration to taskflow. Users may choose to disable/ignore it when unwanted.

Miretpl · 2026-03-04T18:47:35Z

I think there should be a separate rule to assist with migration to taskflow. Users may choose to disable/ignore it when unwanted.

It could be a separate rule for TaskFlow API to make it easier for some teams to migrate from one way to another, but it will not resolve the issue which I've mentioned.

The rules themselves should be Dag implementation-agnostic in the sense that it doesn't matter if Dag was written in a standard way or with TaskFlow API. The different approach would cover only a subset of cases, while both ways are equally supported, and there is no recommended way of Dags writing in that sense (at least I didn't find anything in the documentation regarding it). Looking at the current implementation of task-branch-as-short-circuit rule, I think it should detect both:

functions decorated by task.branch decorator
functions which are used in BranchPythonOperator
and do proper logic execution to detect if the code is following best practices or not.

P.S. I found two notes in the Airflow documentation regarding the recommendation related to the TaskFlow API. One is connected to the task.branch, which is a recommended way, and the other to the task, which is recommended unless there is no template rendering (nothing generally applicable to the TaskFlow itself, but there is a chance that I've missed something).

Dev-iL · 2026-03-04T23:02:51Z

Looking at the current implementation of task-branch-as-short-circuit rule, I think it should detect both:

functions decorated by task.branch decorator

functions which are used in BranchPythonOperator
and do proper logic execution to detect if the code is following best practices or not.

Done in this commit + SKILL updated to be aware of this pattern.

Lee-W · 2026-03-05T02:13:52Z

Hey, thanks for creating this. I'm actaully good with adding this, but it would be better if we had a stronger consensus across the community. I'll start a discussion in the dev list today.

ntBre · 2026-03-05T16:55:01Z

Like #23583, I'll convert this to a draft for now. Thanks very much for the reviews (and the PR!)!

Dev-iL · 2026-03-23T11:30:56Z

@ntBre @Lee-W @sjyangkevin

The Airflow community approved this rule, so I'm marking it "Ready for review".

Refs:

ntBre

Thanks! This looks good to me overall, if @Lee-W is happy with it. I just had a few minor suggestions.

sjyangkevin · 2026-03-29T16:51:59Z

@ntBre @Lee-W @sjyangkevin

The Airflow community approved this rule, so I'm marking it "Ready for review".

Refs:

https://lists.apache.org/thread/odmymy6kn5rqx17x44ydq43c1vfqh511

https://lists.apache.org/thread/q9txg01xvz5gvkjnw2mjsqv1gfqjnbjo

Looks good to me

Lee-W

Looks good to me. Thanks!

Lee-W · 2026-04-02T10:43:02Z

I think we should differentiate different practices so that users can configurate it

ntBre · 2026-04-13T17:15:29Z

Sorry for the delay on review here! Is this ready for another look from me? I wasn't 100% sure that this thread was resolved, but I think it probably is.

Dev-iL · 2026-04-13T17:43:10Z

Sorry for the delay on review here! Is this ready for another look from me? I wasn't 100% sure that this thread was resolved, but I think it probably is.

Yeah, I'd say there's nothing left to discuss there. It was a "meta" discussion anyway. Please take another look!

Lee-W · 2026-04-16T09:39:33Z

yep, I think that one is not for this PR and I don't think we'll do it that way. but we'll need to update the code in the descirption and title

Dev-iL · 2026-04-22T14:53:13Z

@MichaReiser perhaps you can take a look?

ntBre

Thank you!

astral-sh-bot Bot assigned ntBre Feb 26, 2026

sjyangkevin reviewed Feb 26, 2026

View reviewed changes

Dev-iL force-pushed the airflow/branch_short-circuit branch from ce2f212 to 0dec43a Compare February 26, 2026 17:19

Dev-iL mentioned this pull request Feb 28, 2026

[airflow] Extract common utilities for use in new rules #23630

Merged

Dev-iL force-pushed the airflow/branch_short-circuit branch 2 times, most recently from 40466f0 to bee9aef Compare March 4, 2026 07:14

Dev-iL force-pushed the airflow/branch_short-circuit branch from bee9aef to 63ee3d1 Compare March 4, 2026 23:02

Dev-iL force-pushed the airflow/branch_short-circuit branch from 63ee3d1 to 9d964b6 Compare March 5, 2026 07:56

ntBre marked this pull request as draft March 5, 2026 16:55

Dev-iL changed the title ~~[airflow] Implement task-branch-as-short-circuit (AIR003)~~ [airflow] Implement task-branch-as-short-circuit (AIR005) Mar 20, 2026

Dev-iL force-pushed the airflow/branch_short-circuit branch from 9d964b6 to 05cbc24 Compare March 20, 2026 12:09

Dev-iL marked this pull request as ready for review March 23, 2026 11:30

astral-sh-bot Bot requested a review from ntBre March 23, 2026 11:31

ntBre added rule Implementing or modifying a lint rule preview Related to preview mode features labels Mar 24, 2026

ntBre reviewed Mar 24, 2026

View reviewed changes

Dev-iL force-pushed the airflow/branch_short-circuit branch from 05cbc24 to 7076aa3 Compare March 25, 2026 09:23

Lee-W reviewed Apr 2, 2026

View reviewed changes

Dev-iL force-pushed the airflow/branch_short-circuit branch from 7076aa3 to 93bd5e4 Compare April 14, 2026 08:10

Implement the branch-to-short-circuit rule (AIR004)

e01f6e6

Dev-iL force-pushed the airflow/branch_short-circuit branch from 93bd5e4 to e01f6e6 Compare April 16, 2026 09:48

Dev-iL changed the title ~~[airflow] Implement task-branch-as-short-circuit (AIR005)~~ [airflow] Implement task-branch-as-short-circuit (AIR004) Apr 16, 2026

Dev-iL changed the title ~~[airflow] Implement task-branch-as-short-circuit (AIR004)~~ [airflow] Implement task-branch-as-short-circuit (AIR004) Apr 16, 2026

ntBre approved these changes Apr 23, 2026

View reviewed changes

ntBre changed the title ~~[airflow] Implement task-branch-as-short-circuit (AIR004)~~ [airflow] Implement task-branch-as-short-circuit (AIR004) Apr 23, 2026

ntBre merged commit 94e6110 into astral-sh:main Apr 23, 2026
44 checks passed

Dev-iL deleted the airflow/branch_short-circuit branch April 23, 2026 14:57

BrewTestBot mentioned this pull request Apr 24, 2026

ruff 0.15.12 Homebrew/homebrew-core#279335

Merged

This was referenced Apr 27, 2026

Always include panic payload in panic diagnostic message #24873

Merged

ASCII identifier fast path #24876

Draft

Conversation

Dev-iL commented Feb 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation details

What it does NOT flag

Test Plan

Uh oh!

ntBre commented Feb 26, 2026

Uh oh!

Lee-W commented Feb 26, 2026

Uh oh!

sjyangkevin left a comment

Choose a reason for hiding this comment

Uh oh!

sjyangkevin Feb 26, 2026

Choose a reason for hiding this comment

Uh oh!

Dev-iL Feb 26, 2026

Choose a reason for hiding this comment

Uh oh!

sjyangkevin Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Dev-iL Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Dev-iL Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Lee-W Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

astral-sh-bot Bot commented Feb 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

ruff-ecosystem results

Linter (stable)

Linter (preview)

Uh oh!

Miretpl commented Mar 3, 2026

Uh oh!

Dev-iL commented Mar 3, 2026

Uh oh!

Miretpl commented Mar 3, 2026

Uh oh!

Dev-iL commented Mar 4, 2026

Uh oh!

Miretpl commented Mar 4, 2026

Uh oh!

Dev-iL commented Mar 4, 2026

Uh oh!

Lee-W commented Mar 5, 2026

Uh oh!

ntBre commented Mar 5, 2026

Uh oh!

Dev-iL commented Mar 23, 2026

Uh oh!

ntBre left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sjyangkevin commented Mar 29, 2026

Uh oh!

Lee-W left a comment

Choose a reason for hiding this comment

Uh oh!

Lee-W Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

ntBre commented Apr 13, 2026

Uh oh!

Dev-iL commented Apr 13, 2026

Uh oh!

Lee-W commented Apr 16, 2026

Uh oh!

Dev-iL commented Feb 26, 2026 •

edited

Loading

astral-sh-bot Bot commented Feb 26, 2026 •

edited

Loading

`ruff-ecosystem` results