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 7e0e11d065..3dd06cb0a7 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 @@ -157,7 +157,7 @@ const REMOVED_CONTEXT_KEYS: [&str; 12] = [ /// pass /// ``` fn check_function_parameters(checker: &mut Checker, function_def: &StmtFunctionDef) { - if !is_airflow_task_function_def(function_def, checker.semantic()) + if !is_airflow_task(function_def, checker.semantic()) && !is_execute_method_inherits_from_airflow_operator(function_def, checker.semantic()) { return; @@ -351,7 +351,7 @@ fn check_class_attribute(checker: &mut Checker, attribute_expr: &ExprAttribute) /// context.get("conf") # 'conf' is removed in Airflow 3.0 /// ``` fn check_context_key_usage_in_call(checker: &mut Checker, call_expr: &ExprCall) { - if !in_airflow_task_function_def(checker.semantic()) { + if !in_airflow_task_function(checker.semantic()) { return; } @@ -400,7 +400,7 @@ fn check_context_key_usage_in_call(checker: &mut Checker, call_expr: &ExprCall) /// Check if a subscript expression accesses a removed Airflow context variable. /// If a removed key is found, push a corresponding diagnostic. fn check_context_key_usage_in_subscript(checker: &mut Checker, subscript: &ExprSubscript) { - if !in_airflow_task_function_def(checker.semantic()) { + if !in_airflow_task_function(checker.semantic()) { return; } @@ -1064,15 +1064,15 @@ fn is_airflow_builtin_or_provider(segments: &[&str], module: &str, symbol_suffix /// Returns `true` if the current statement hierarchy has a function that's decorated with /// `@airflow.decorators.task`. -fn in_airflow_task_function_def(semantic: &SemanticModel) -> bool { +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(function_def, semantic)) + .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(function_def: &StmtFunctionDef, semantic: &SemanticModel) -> bool { +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))