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

Co-authored-by: Claude Opus 4.6 <[email protected]>

Author: Dev-iL

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

@@ -205,3 +205,37 @@ def test_inlet_events_dataset_subscript_ok(**context):
 
     print(context["inlet_events"][Dataset("this://is-url")])
     print(context["inlet_events"][Asset("this://is-url")])
+
+
+# Same context checks with airflow.sdk import
+from airflow.sdk import task as sdk_task
+
+
+@sdk_task
+def sdk_access_deprecated_context_key(**context):
+    execution_date = context["execution_date"]
+    next_ds = context["next_ds"]
+
+
+@sdk_task
+def sdk_access_valid_context_key(**context):
+    logical_date = context["logical_date"]
+
+
+# Test variant decorator forms like @task.branch and @task.short_circuit
+@task.branch
+def branch_task_with_deprecated_context(**context):
+    execution_date = context["execution_date"]
+    return "some_task"
+
+
+@task.short_circuit
+def short_circuit_task_with_deprecated_context(**context):
+    next_ds = context["next_ds"]
+    return True
+
+
+@task.branch()
+def branch_task_with_call_and_deprecated_context(**context):
+    tomorrow_ds = context["tomorrow_ds"]
+    return "some_task"

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

@@ -3,6 +3,7 @@ use crate::fix::edits::remove_unused_imports;
 use crate::importer::ImportRequest;
 use crate::rules::numpy::helpers::{AttributeSearcher, ImportSearcher};
 use ruff_diagnostics::{Edit, Fix};
+use ruff_python_ast::helpers::map_callable;
 use ruff_python_ast::name::{QualifiedName, QualifiedNameBuilder};
 use ruff_python_ast::statement_visitor::StatementVisitor;
 use ruff_python_ast::visitor::Visitor;
@@ -290,3 +291,38 @@ where
 
     any_qualified_base_class(class_def, semantic, &is_base_class)
 }
+
+/// Returns `true` if the current statement hierarchy has a function that's decorated with
+/// `@airflow.decorators.task` or `@airflow.sdk.task`.
+pub(crate) fn in_airflow_task_function(semantic: &SemanticModel) -> bool {
+    semantic
+        .current_statements()
+        .find_map(|stmt| stmt.as_function_def_stmt())
+        .is_some_and(|function_def| is_airflow_task(function_def, semantic))
+}
+
+/// Returns `true` if the given function is decorated with `@airflow.decorators.task`
+/// (or `@airflow.sdk.task`), including variant forms like `@task.branch` and
+/// `@task.short_circuit`.
+pub(crate) fn is_airflow_task(function_def: &StmtFunctionDef, semantic: &SemanticModel) -> bool {
+    function_def.decorator_list.iter().any(|decorator| {
+        let expr = map_callable(&decorator.expression);
+
+        // Match `@task` and `@task()` directly.
+        if semantic
+            .resolve_qualified_name(expr)
+            .is_some_and(|qn| matches!(qn.segments(), ["airflow", "decorators" | "sdk", "task"]))
+        {
+            return true;
+        }
+
+        // Match `@task.<variant>` (e.g., `@task.branch`, `@task.short_circuit`).
+        if let Expr::Attribute(ExprAttribute { value, .. }) = expr {
+            return semantic.resolve_qualified_name(value).is_some_and(|qn| {
+                matches!(qn.segments(), ["airflow", "decorators" | "sdk", "task"])
+            });
+        }
+
+        false
+    })
+}

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

@@ -1,7 +1,8 @@
 use crate::checkers::ast::Checker;
 use crate::rules::airflow::helpers::{
     Replacement, generate_import_edit, generate_remove_and_runtime_import_edit,
-    is_airflow_builtin_or_provider, is_guarded_by_try_except, is_method_in_subclass,
+    in_airflow_task_function, is_airflow_builtin_or_provider, is_airflow_task,
+    is_guarded_by_try_except, is_method_in_subclass,
 };
 use crate::{Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{ViolationMetadata, derive_message_formats};
@@ -1223,26 +1224,6 @@ fn is_airflow_auth_manager(segments: &[&str]) -> bool {
     }
 }
 
-/// Returns `true` if the current statement hierarchy has a function that's decorated with
-/// `@airflow.decorators.task`.
-fn in_airflow_task_function(semantic: &SemanticModel) -> bool {
-    semantic
-        .current_statements()
-        .find_map(|stmt| stmt.as_function_def_stmt())
-        .is_some_and(|function_def| is_airflow_task(function_def, semantic))
-}
-
-/// Returns `true` if the given function is decorated with `@airflow.decorators.task`.
-fn is_airflow_task(function_def: &StmtFunctionDef, semantic: &SemanticModel) -> bool {
-    function_def.decorator_list.iter().any(|decorator| {
-        semantic
-            .resolve_qualified_name(map_callable(&decorator.expression))
-            .is_some_and(|qualified_name| {
-                matches!(qualified_name.segments(), ["airflow", "decorators", "task"])
-            })
-    })
-}
-
 /// Check it's "execute" method inherits from Airflow base operator
 ///
 /// For example:

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR301_AIR301_context.py.snap

@@ -521,3 +521,52 @@ AIR301 `inlet_events["<uri>"]` is removed in Airflow 3.0
     |                            ^^^^^^^^^^^^^^^
     |
 help: Accessing `inlet_events` via a string key is deprecated; use `context["inlet_events"][Asset(uri="this://is-url")]` instead of `context["inlet_events"]["this://is-url"]`.
+
+AIR301 `execution_date` is removed in Airflow 3.0
+   --> AIR301_context.py:216:30
+    |
+214 | @sdk_task
+215 | def sdk_access_deprecated_context_key(**context):
+216 |     execution_date = context["execution_date"]
+    |                              ^^^^^^^^^^^^^^^^
+217 |     next_ds = context["next_ds"]
+    |
+
+AIR301 `next_ds` is removed in Airflow 3.0
+   --> AIR301_context.py:217:23
+    |
+215 | def sdk_access_deprecated_context_key(**context):
+216 |     execution_date = context["execution_date"]
+217 |     next_ds = context["next_ds"]
+    |                       ^^^^^^^^^
+    |
+
+AIR301 `execution_date` is removed in Airflow 3.0
+   --> AIR301_context.py:228:30
+    |
+226 | @task.branch
+227 | def branch_task_with_deprecated_context(**context):
+228 |     execution_date = context["execution_date"]
+    |                              ^^^^^^^^^^^^^^^^
+229 |     return "some_task"
+    |
+
+AIR301 `next_ds` is removed in Airflow 3.0
+   --> AIR301_context.py:234:23
+    |
+232 | @task.short_circuit
+233 | def short_circuit_task_with_deprecated_context(**context):
+234 |     next_ds = context["next_ds"]
+    |                       ^^^^^^^^^
+235 |     return True
+    |
+
+AIR301 `tomorrow_ds` is removed in Airflow 3.0
+   --> AIR301_context.py:240:27
+    |
+238 | @task.branch()
+239 | def branch_task_with_call_and_deprecated_context(**context):
+240 |     tomorrow_ds = context["tomorrow_ds"]
+    |                           ^^^^^^^^^^^^^
+241 |     return "some_task"
+    |