From 8a7ec4c0a30e28097073c93a41015353ea7d142e Mon Sep 17 00:00:00 2001 From: Ankit Chaurasia <8670962+sunank200@users.noreply.github.com> Date: Wed, 15 Jan 2025 12:20:02 +0545 Subject: [PATCH] Fix PR comments --- .../src/rules/airflow/rules/removal_in_3.rs | 59 ++-- ...flow__tests__AIR302_AIR302_context.py.snap | 329 ++++++++++++++++++ 2 files changed, 363 insertions(+), 25 deletions(-) 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 83c889b884..ea5c30ca13 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 @@ -319,18 +319,15 @@ fn check_class_attribute(checker: &mut Checker, attribute_expr: &ExprAttribute) /// print("access invalid key", context.get("conf")) /// ``` fn check_context_get(checker: &mut Checker, call_expr: &ExprCall) { - if is_task_context_referenced(checker, &call_expr.func) { - return; - } + add_context_diagnostics(checker, &call_expr.func); let Expr::Attribute(ExprAttribute { value, attr, .. }) = &*call_expr.func else { return; }; - if !value - .as_name_expr() - .is_some_and(|name| matches!(name.id.as_str(), "context" | "kwargs")) - { + if !value.as_name_expr().is_some_and(|name| { + matches!(name.id.as_str(), "context" | "kwargs") || name.id.as_str().starts_with("**") + }) { return; } @@ -961,33 +958,45 @@ fn is_airflow_builtin_or_provider(segments: &[&str], module: &str, symbol_suffix } } -fn is_task_context_referenced(checker: &mut Checker, expr: &Expr) -> bool { +fn uses_removed_context_keys(checker: &mut Checker) -> bool { let parents: Vec<_> = checker.semantic().current_statements().collect(); for stmt in parents { - if let Stmt::FunctionDef(function_def) = stmt { - if is_decorated_with(checker, function_def) { - let arguments = extract_task_function_arguments(function_def); + let Stmt::FunctionDef(function_def) = stmt else { + continue; + }; - for deprecated_arg in REMOVED_CONTEXT_KEYS { - if arguments.contains(&deprecated_arg.to_string()) { - checker.diagnostics.push(Diagnostic::new( - Airflow3Removal { - deprecated: deprecated_arg.to_string(), - replacement: Replacement::None, - }, - expr.range(), - )); - return true; - } - } - } + if !has_task_decorator(checker, function_def) { + continue; + } + + let arguments = extract_task_function_arguments(function_def); + + if arguments + .iter() + .any(|arg| REMOVED_CONTEXT_KEYS.contains(&arg.as_str())) + { + return true; } } false } +fn add_context_diagnostics(checker: &mut Checker, expr: &Expr) { + if uses_removed_context_keys(checker) { + for removed_key in REMOVED_CONTEXT_KEYS { + checker.diagnostics.push(Diagnostic::new( + Airflow3Removal { + deprecated: removed_key.to_string(), + replacement: Replacement::None, + }, + expr.range(), + )); + } + } +} + fn extract_task_function_arguments(stmt: &StmtFunctionDef) -> Vec { let mut arguments = Vec::new(); @@ -1002,7 +1011,7 @@ fn extract_task_function_arguments(stmt: &StmtFunctionDef) -> Vec { arguments } -fn is_decorated_with(checker: &mut Checker, stmt: &StmtFunctionDef) -> bool { +fn has_task_decorator(checker: &mut Checker, stmt: &StmtFunctionDef) -> bool { stmt.decorator_list.iter().any(|decorator| { checker .semantic() diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap index 57fb660978..731230dbcd 100644 --- a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap @@ -404,6 +404,15 @@ AIR302_context.py:111:39: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0 113 | @task | +AIR302_context.py:115:5: AIR302 `conf` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + AIR302_context.py:115:5: AIR302 `execution_date` is removed in Airflow 3.0 | 113 | @task @@ -413,6 +422,106 @@ AIR302_context.py:115:5: AIR302 `execution_date` is removed in Airflow 3.0 116 | print("access invalid key", context.get("conf")) | +AIR302_context.py:115:5: AIR302 `next_ds` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `next_ds_nodash` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `next_execution_date` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `prev_ds` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `prev_ds_nodash` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `prev_execution_date` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `prev_execution_date_success` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `tomorrow_ds` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `yesterday_ds` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:115:5: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0 + | +113 | @task +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) + | ^^^^^ AIR302 +116 | print("access invalid key", context.get("conf")) + | + +AIR302_context.py:116:5: AIR302 `conf` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + AIR302_context.py:116:5: AIR302 `execution_date` is removed in Airflow 3.0 | 114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): @@ -423,6 +532,116 @@ AIR302_context.py:116:5: AIR302 `execution_date` is removed in Airflow 3.0 118 | @task(task_id="print_the_context") | +AIR302_context.py:116:5: AIR302 `next_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `next_ds_nodash` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `next_execution_date` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `prev_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `prev_ds_nodash` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `prev_execution_date` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `prev_execution_date_success` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `tomorrow_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `yesterday_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:5: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `conf` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + AIR302_context.py:116:33: AIR302 `execution_date` is removed in Airflow 3.0 | 114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): @@ -433,6 +652,116 @@ AIR302_context.py:116:33: AIR302 `execution_date` is removed in Airflow 3.0 118 | @task(task_id="print_the_context") | +AIR302_context.py:116:33: AIR302 `next_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `next_ds_nodash` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `next_execution_date` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `prev_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `prev_ds_nodash` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `prev_execution_date` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `prev_execution_date_success` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `tomorrow_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `yesterday_ds` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:33: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + +AIR302_context.py:116:45: AIR302 `conf` is removed in Airflow 3.0 + | +114 | def access_invalid_argument_task_out_of_dag(execution_date, **context): +115 | print("execution date", execution_date) +116 | print("access invalid key", context.get("conf")) + | ^^^^^^ AIR302 +117 | +118 | @task(task_id="print_the_context") + | + AIR302_context.py:122:22: AIR302 `conf` is removed in Airflow 3.0 | 120 | """Print the Airflow context and ds variable from the context."""