feat(AIR302): add is_execute_method_inherits_from_airflow_operator for checking removed context key in execute method

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

@@ -154,7 +154,9 @@ const REMOVED_CONTEXT_KEYS: [&str; 12] = [
 ///     pass
 /// ```
 fn check_function_parameters(checker: &mut Checker, function_def: &StmtFunctionDef) {
-    if !is_airflow_task(function_def, checker.semantic()) {
+    if !is_airflow_task(function_def, checker.semantic())
+        && !is_execute_method_inherits_from_airflow_operator(function_def, checker.semantic())
+    {
         return;
     }
 
@@ -1076,3 +1078,39 @@ fn is_airflow_task(function_def: &StmtFunctionDef, semantic: &SemanticModel) ->
             })
     })
 }
+
+/// Check it's "execute" method inherits from Airflow base operator
+///
+/// For example:
+///
+/// ```python
+/// from airflow.models.baseoperator import BaseOperator
+///
+/// class CustomOperator(BaseOperator):
+///     def execute(self):
+///         pass
+/// ```
+fn is_execute_method_inherits_from_airflow_operator(
+    function_def: &StmtFunctionDef,
+    semantic: &SemanticModel,
+) -> bool {
+    if function_def.name.as_str() != "execute" {
+        return false;
+    }
+
+    let ScopeKind::Class(class_def) = semantic.current_scope().kind else {
+        return false;
+    };
+
+    if !class_def.bases().iter().any(|class_base| {
+        semantic
+            .resolve_qualified_name(class_base)
+            .is_some_and(|qualified_name| {
+                matches!(qualified_name.segments(), ["airflow", .., "BaseOperator"])
+            })
+    }) {
+        return false;
+    };
+
+    true
+}

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

@@ -327,3 +327,12 @@ AIR302_context.py:115:13: AIR302 `airflow.operators.dummy.DummyOperator` is remo
 117 |         params={
     |
     = help: Use `airflow.operators.empty.EmptyOperator` instead
+
+AIR302_context.py:135:23: AIR302 `next_ds` is removed in Airflow 3.0
+    |
+134 | class CustomOperator(BaseOperator):
+135 |     def execute(self, next_ds, context):
+    |                       ^^^^^^^ AIR302
+136 |         execution_date = context["execution_date"]
+137 |         next_ds = context["next_ds"]
+    |