From e8445681bc5e8e8234db6075edb707b29f265773 Mon Sep 17 00:00:00 2001
From: Ankit Chaurasia <8670962+sunank200@users.noreply.github.com>
Date: Thu, 23 Jan 2025 14:46:13 +0545
Subject: [PATCH]  use a separate entrypoint which directly checks the function
 definition

---
 .../src/checkers/ast/analyze/statement.rs     |  3 +
 .../src/rules/airflow/rules/removal_in_3.rs   | 92 ++++++++++---------
 2 files changed, 54 insertions(+), 41 deletions(-)

diff --git a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs
index d5f92b9f293f2..7f64be4bdaaad 100644
--- a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs
+++ b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs
@@ -376,6 +376,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::PytestParameterWithDefaultArgument) {
                 flake8_pytest_style::rules::parameter_with_default_argument(checker, function_def);
             }
+            if checker.enabled(Rule::Airflow3Removal) {
+                airflow::rules::removed_in_3_function_def(checker, function_def);
+            }
         }
         Stmt::Return(_) => {
             if checker.enabled(Rule::ReturnOutsideFunction) {
diff --git a/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs b/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs
index 3a2f9a6f638a5..c356ae706618e 100644
--- a/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs
+++ b/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs
@@ -342,16 +342,6 @@ fn find_parameter<'a>(
 ///     print(context.get("conf"))
 /// ```
 ///
-/// **Removed context variable as a parameter:**
-/// ```python
-/// from airflow.decorators import task
-///
-/// @task
-/// def another_task(execution_date, **kwargs):
-///     # 'execution_date' is removed in Airflow 3.0
-///     pass
-/// ```
-///
 /// **Accessing multiple keys:**
 /// ```python
 /// from airflow.decorators import task
@@ -398,37 +388,6 @@ fn check_removed_context_keys_usage(checker: &mut Checker, call_expr: &ExprCall)
     if attr.as_str() != "get" {
         return;
     }
-    let function_def = {
-        let mut parents = checker.semantic().current_statements();
-        parents.find_map(|stmt| {
-            if let Stmt::FunctionDef(func_def) = stmt {
-                Some(func_def.clone())
-            } else {
-                None
-            }
-        })
-    };
-
-    if let Some(func_def) = function_def {
-        for param in func_def
-            .parameters
-            .posonlyargs
-            .iter()
-            .chain(func_def.parameters.args.iter())
-            .chain(func_def.parameters.kwonlyargs.iter())
-        {
-            let param_name = param.parameter.name.as_str();
-            if REMOVED_CONTEXT_KEYS.contains(&param_name) {
-                checker.diagnostics.push(Diagnostic::new(
-                    Airflow3Removal {
-                        deprecated: param_name.to_string(),
-                        replacement: Replacement::None,
-                    },
-                    param.parameter.name.range(),
-                ));
-            }
-        }
-    }
 
     for removed_key in REMOVED_CONTEXT_KEYS {
         if let Some(argument) = call_expr.arguments.find_argument_value(removed_key, 0) {
@@ -1087,3 +1046,54 @@ fn is_airflow_builtin_or_provider(segments: &[&str], module: &str, symbol_suffix
         _ => false,
     }
 }
+
+/// AIR302 Check the function argument for removed context variable.
+/// For example:
+/// **Removed context variable as a parameter:**
+/// ```python
+/// from airflow.decorators import task
+///
+/// @task
+/// def another_task(execution_date, **kwargs):
+///     # 'execution_date' is removed in Airflow 3.0
+///     pass
+/// ```
+pub(crate) fn removed_in_3_function_def(checker: &mut Checker, function_def: &StmtFunctionDef) {
+    if !checker.semantic().seen_module(Modules::AIRFLOW) {
+        return;
+    }
+
+    if !is_airflow_task(function_def, checker.semantic()) {
+        return;
+    }
+
+    for param in function_def
+        .parameters
+        .posonlyargs
+        .iter()
+        .chain(function_def.parameters.args.iter())
+        .chain(function_def.parameters.kwonlyargs.iter())
+    {
+        let param_name = param.parameter.name.as_str();
+        if REMOVED_CONTEXT_KEYS.contains(&param_name) {
+            checker.diagnostics.push(Diagnostic::new(
+                Airflow3Removal {
+                    deprecated: param_name.to_string(),
+                    replacement: Replacement::None,
+                },
+                param.parameter.name.range(),
+            ));
+        }
+    }
+}
+
+/// 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"])
+            })
+    })
+}