Add more checks for removed context variables

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

Author: sunank200

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

@@ -0,0 +1,76 @@
+from airflow.models import DAG
+from airflow.operators.dummy import DummyOperator
+from datetime import datetime
+from airflow.plugins_manager import AirflowPlugin
+from airflow.decorators import task, get_current_context
+from airflow.models.baseoperator import BaseOperator
+
+@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"]
+
+with DAG(
+    dag_id="example_dag",
+    schedule_interval="@daily",
+    start_date=datetime(2023, 1, 1),
+    template_searchpath=["/templates"],
+) as dag:
+    task1 = DummyOperator(
+        task_id="task1",
+        params={
+            # Removed variables in template
+            "execution_date": "{{ execution_date }}",
+            "next_ds": "{{ next_ds }}",
+            "prev_ds": "{{ prev_ds }}"
+        },
+    )
+
+class CustomMacrosPlugin(AirflowPlugin):
+    name = "custom_macros"
+    macros = {
+        "execution_date_macro": lambda context: context["execution_date"],
+        "next_ds_macro": lambda context: context["next_ds"]
+    }
+
+@task
+def print_config():
+    context = get_current_context()
+    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"]
+
+class CustomOperator(BaseOperator):
+    def execute(self, context):
+        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

@@ -18,6 +18,7 @@ mod tests {
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_class_attribute.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_airflow_plugin.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

@@ -2,7 +2,7 @@ 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, ExprContext, ExprName,
-    StmtClassDef,
+    ExprStringLiteral, ExprSubscript, StmtClassDef,
 };
 use ruff_python_semantic::analyze::typing;
 use ruff_python_semantic::Modules;
@@ -71,6 +71,63 @@ impl Violation for Airflow3Removal {
     }
 }
 
+fn extract_name_from_slice(slice: &Expr) -> Option<String> {
+    match slice {
+        Expr::StringLiteral(ExprStringLiteral { value, .. }) => Some(value.to_string()),
+        _ => None,
+    }
+}
+
+pub(crate) fn removed_context_variable(checker: &mut Checker, expr: &Expr) {
+    const REMOVED_CONTEXT_KEYS: [&str; 12] = [
+        "conf",
+        "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 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) {
+                    if REMOVED_CONTEXT_KEYS.contains(&key.as_str()) {
+                        checker.diagnostics.push(Diagnostic::new(
+                            Airflow3Removal {
+                                deprecated: key,
+                                replacement: Replacement::None,
+                            },
+                            slice.range(),
+                        ));
+                    }
+                }
+            }
+        }
+    }
+
+    if let Expr::StringLiteral(ExprStringLiteral { value, .. }) = expr {
+        let value_str = value.to_string();
+        for key in REMOVED_CONTEXT_KEYS {
+            if value_str.contains(&format!("{{{{ {key} }}}}")) {
+                checker.diagnostics.push(Diagnostic::new(
+                    Airflow3Removal {
+                        deprecated: key.to_string(),
+                        replacement: Replacement::None,
+                    },
+                    expr.range(),
+                ));
+            }
+        }
+    }
+}
+
 /// AIR302
 pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
     if !checker.semantic().seen_module(Modules::AIRFLOW) {
@@ -100,6 +157,9 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
                 }
             }
         }
+        Expr::Subscript(_) => {
+            removed_context_variable(checker, expr);
+        }
         _ => {}
     }
 }