add lint rule to show error for removed context variables in airflow

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

@@ -0,0 +1,19 @@
+from airflow.decorators import task
+
+@task
+def print_config(**context):
+    # This should not throw an error as logical_date is part of airflow context.
+    logical_date = context["logical_date"]
+
+    # Removed usage - should trigger violations
+    execution_date = context["execution_date"]
+    next_ds = context["next_ds"]
+    next_ds_nodash = context["next_ds_nodash"]
+    next_execution_date = context["next_execution_date"]
+    prev_ds = context["prev_ds"]
+    prev_ds_nodash = context["prev_ds_nodash"]
+    prev_execution_date = context["prev_execution_date"]
+    prev_execution_date_success = context["prev_execution_date_success"]
+    tomorrow_ds = context["tomorrow_ds"]
+    yesterday_ds = context["yesterday_ds"]
+    yesterday_ds_nodash = context["yesterday_ds_nodash"]

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

@@ -171,7 +171,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::NonPEP646Unpack) {
                 pyupgrade::rules::use_pep646_unpack(checker, subscript);
             }
-
+            if checker.enabled(Rule::Airflow3Removal) {
+                airflow::rules::removed_in_3(checker, expr);
+            }
             pandas_vet::rules::subscript(checker, value, expr);
         }
         Expr::Tuple(ast::ExprTuple {

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

@@ -16,6 +16,7 @@ mod tests {
     #[test_case(Rule::AirflowDagNoScheduleArgument, Path::new("AIR301.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_args.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))]
+    #[test_case(Rule::Airflow3Removal, Path::new("AIR302_context.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR303.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());

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

@@ -1,6 +1,9 @@
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::{name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall};
+use ruff_python_ast::{
+    name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall, ExprName, ExprStringLiteral,
+    ExprSubscript,
+};
 use ruff_python_semantic::analyze::typing;
 use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
@@ -41,6 +44,7 @@ enum Replacement {
 pub(crate) struct Airflow3Removal {
     deprecated: String,
     replacement: Replacement,
+    is_context_variable: bool,
 }
 
 impl Violation for Airflow3Removal {
@@ -51,14 +55,19 @@ impl Violation for Airflow3Removal {
         let Airflow3Removal {
             deprecated,
             replacement,
+            is_context_variable,
         } = self;
         match replacement {
             Replacement::None => format!("`{deprecated}` is removed in Airflow 3.0"),
             Replacement::Name(_) => {
                 format!("`{deprecated}` is removed in Airflow 3.0")
             }
             Replacement::Message(message) => {
-                format!("`{deprecated}` is removed in Airflow 3.0; {message}")
+                if *is_context_variable {
+                    format!("`{deprecated}` {message}")
+                } else {
+                    format!("`{deprecated}` is removed in Airflow 3.0; {message}")
+                }
             }
         }
     }
@@ -86,6 +95,7 @@ fn diagnostic_for_argument(
                 Some(name) => Replacement::Name(name),
                 None => Replacement::None,
             },
+            is_context_variable: false,
         },
         keyword
             .arg
@@ -203,6 +213,7 @@ fn removed_method(checker: &mut Checker, expr: &Expr) {
             Airflow3Removal {
                 deprecated: attr.to_string(),
                 replacement,
+                is_context_variable: false,
             },
             attr.range(),
         ));
@@ -640,12 +651,57 @@ fn removed_name(checker: &mut Checker, expr: &Expr, ranged: impl Ranged) {
             Airflow3Removal {
                 deprecated,
                 replacement,
+                is_context_variable: false,
             },
             ranged.range(),
         ));
     }
 }
 
+fn extract_name_from_slice(slice: &Expr) -> Option<String> {
+    if let Expr::StringLiteral(ExprStringLiteral { value, .. }) = slice {
+        Some(value.to_string())
+    } else {
+        None
+    }
+}
+
+pub(crate) fn removed_context_variable(checker: &mut Checker, expr: &Expr) {
+    if let Expr::Subscript(ExprSubscript { value, slice, .. }) = expr {
+        if let Expr::Name(ExprName { id, .. }) = &**value {
+            if id.as_str() == "context" {
+                if let Some(key) = extract_name_from_slice(slice) {
+                    const REMOVED_CONTEXT_KEYS: [&str; 11] = [
+                        "execution_date",
+                        "next_ds",
+                        "next_ds_nodash",
+                        "next_execution_date",
+                        "prev_ds",
+                        "prev_ds_nodash",
+                        "prev_execution_date",
+                        "prev_execution_date_success",
+                        "tomorrow_ds",
+                        "yesterday_ds",
+                        "yesterday_ds_nodash",
+                    ];
+                    if REMOVED_CONTEXT_KEYS.contains(&key.as_str()) {
+                        checker.diagnostics.push(Diagnostic::new(
+                            Airflow3Removal {
+                                deprecated: key,
+                                replacement: Replacement::Message(
+                                    "is removed in the Airflow context in Airflow 3.0",
+                                ),
+                                is_context_variable: true,
+                            },
+                            slice.range(),
+                        ));
+                    }
+                }
+            }
+        }
+    }
+}
+
 /// AIR302
 pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
     if !checker.semantic().seen_module(Modules::AIRFLOW) {
@@ -662,8 +718,13 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
 
             removed_method(checker, expr);
         }
-        Expr::Attribute(ExprAttribute { attr: ranged, .. }) => removed_name(checker, expr, ranged),
+        Expr::Attribute(ExprAttribute { attr, .. }) => {
+            removed_name(checker, expr, attr);
+        }
         ranged @ Expr::Name(_) => removed_name(checker, expr, ranged),
+        Expr::Subscript(_) => {
+            removed_context_variable(checker, expr);
+        }
         _ => {}
     }
 }

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

@@ -0,0 +1,109 @@
+---
+source: crates/ruff_linter/src/rules/airflow/mod.rs
+snapshot_kind: text
+---
+AIR302_context.py:9:30: AIR302 `execution_date` is removed in the Airflow context in Airflow 3.0
+   |
+ 8 |     # Removed usage - should trigger violations
+ 9 |     execution_date = context["execution_date"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+10 |     next_ds = context["next_ds"]
+11 |     next_ds_nodash = context["next_ds_nodash"]
+   |
+
+AIR302_context.py:10:23: AIR302 `next_ds` is removed in the Airflow context in Airflow 3.0
+   |
+ 8 |     # Removed usage - should trigger violations
+ 9 |     execution_date = context["execution_date"]
+10 |     next_ds = context["next_ds"]
+   |                       ^^^^^^^^^ AIR302
+11 |     next_ds_nodash = context["next_ds_nodash"]
+12 |     next_execution_date = context["next_execution_date"]
+   |
+
+AIR302_context.py:11:30: AIR302 `next_ds_nodash` is removed in the Airflow context in Airflow 3.0
+   |
+ 9 |     execution_date = context["execution_date"]
+10 |     next_ds = context["next_ds"]
+11 |     next_ds_nodash = context["next_ds_nodash"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+12 |     next_execution_date = context["next_execution_date"]
+13 |     prev_ds = context["prev_ds"]
+   |
+
+AIR302_context.py:12:35: AIR302 `next_execution_date` is removed in the Airflow context in Airflow 3.0
+   |
+10 |     next_ds = context["next_ds"]
+11 |     next_ds_nodash = context["next_ds_nodash"]
+12 |     next_execution_date = context["next_execution_date"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+13 |     prev_ds = context["prev_ds"]
+14 |     prev_ds_nodash = context["prev_ds_nodash"]
+   |
+
+AIR302_context.py:13:23: AIR302 `prev_ds` is removed in the Airflow context in Airflow 3.0
+   |
+11 |     next_ds_nodash = context["next_ds_nodash"]
+12 |     next_execution_date = context["next_execution_date"]
+13 |     prev_ds = context["prev_ds"]
+   |                       ^^^^^^^^^ AIR302
+14 |     prev_ds_nodash = context["prev_ds_nodash"]
+15 |     prev_execution_date = context["prev_execution_date"]
+   |
+
+AIR302_context.py:14:30: AIR302 `prev_ds_nodash` is removed in the Airflow context in Airflow 3.0
+   |
+12 |     next_execution_date = context["next_execution_date"]
+13 |     prev_ds = context["prev_ds"]
+14 |     prev_ds_nodash = context["prev_ds_nodash"]
+   |                              ^^^^^^^^^^^^^^^^ AIR302
+15 |     prev_execution_date = context["prev_execution_date"]
+16 |     prev_execution_date_success = context["prev_execution_date_success"]
+   |
+
+AIR302_context.py:15:35: AIR302 `prev_execution_date` is removed in the Airflow context in Airflow 3.0
+   |
+13 |     prev_ds = context["prev_ds"]
+14 |     prev_ds_nodash = context["prev_ds_nodash"]
+15 |     prev_execution_date = context["prev_execution_date"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+16 |     prev_execution_date_success = context["prev_execution_date_success"]
+17 |     tomorrow_ds = context["tomorrow_ds"]
+   |
+
+AIR302_context.py:16:43: AIR302 `prev_execution_date_success` is removed in the Airflow context in Airflow 3.0
+   |
+14 |     prev_ds_nodash = context["prev_ds_nodash"]
+15 |     prev_execution_date = context["prev_execution_date"]
+16 |     prev_execution_date_success = context["prev_execution_date_success"]
+   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
+17 |     tomorrow_ds = context["tomorrow_ds"]
+18 |     yesterday_ds = context["yesterday_ds"]
+   |
+
+AIR302_context.py:17:27: AIR302 `tomorrow_ds` is removed in the Airflow context in Airflow 3.0
+   |
+15 |     prev_execution_date = context["prev_execution_date"]
+16 |     prev_execution_date_success = context["prev_execution_date_success"]
+17 |     tomorrow_ds = context["tomorrow_ds"]
+   |                           ^^^^^^^^^^^^^ AIR302
+18 |     yesterday_ds = context["yesterday_ds"]
+19 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |
+
+AIR302_context.py:18:28: AIR302 `yesterday_ds` is removed in the Airflow context in Airflow 3.0
+   |
+16 |     prev_execution_date_success = context["prev_execution_date_success"]
+17 |     tomorrow_ds = context["tomorrow_ds"]
+18 |     yesterday_ds = context["yesterday_ds"]
+   |                            ^^^^^^^^^^^^^^ AIR302
+19 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |
+
+AIR302_context.py:19:35: AIR302 `yesterday_ds_nodash` is removed in the Airflow context in Airflow 3.0
+   |
+17 |     tomorrow_ds = context["tomorrow_ds"]
+18 |     yesterday_ds = context["yesterday_ds"]
+19 |     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+   |                                   ^^^^^^^^^^^^^^^^^^^^^ AIR302
+   |