From f989c2c3af03249a2e564626f2cd1fb0cd04b928 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 2 Apr 2025 23:27:51 +0800 Subject: [PATCH] [`airflow`] Add autofix infrastructure to `AIR302` name checks (#16965) ## Summary Add autofix infrastructure to `AIR302` name checks and use this logic to fix`"airflow", "api_connexion", "security", "requires_access_dataset"`, `"airflow", "Dataset"` and `"airflow", "datasets", "Dataset"` ## Test Plan The existing test fixture reflects the update --- .../test/fixtures/airflow/AIR302_names_try.py | 8 +++ .../ruff_linter/src/rules/airflow/helpers.rs | 64 +++++++++++++++++++ crates/ruff_linter/src/rules/airflow/mod.rs | 2 + .../src/rules/airflow/rules/removal_in_3.rs | 60 +++++++++++------ ...sts__AIR302_AIR302_class_attribute.py.snap | 35 +++++++++- ...irflow__tests__AIR302_AIR302_names.py.snap | 57 ++++++++++++++++- ...ow__tests__AIR302_AIR302_names_try.py.snap | 5 ++ crates/ruff_linter/src/rules/numpy/helpers.rs | 45 +++++++++++++ crates/ruff_linter/src/rules/numpy/mod.rs | 1 + .../numpy/rules/numpy_2_0_deprecation.rs | 47 +------------- 10 files changed, 254 insertions(+), 70 deletions(-) create mode 100644 crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names_try.py create mode 100644 crates/ruff_linter/src/rules/airflow/helpers.rs create mode 100644 crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names_try.py.snap create mode 100644 crates/ruff_linter/src/rules/numpy/helpers.rs diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names_try.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names_try.py new file mode 100644 index 0000000000..5d4c3da974 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names_try.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +try: + from airflow.sdk import Asset +except ModuleNotFoundError: + from airflow.datasets import Dataset as Asset + +Asset diff --git a/crates/ruff_linter/src/rules/airflow/helpers.rs b/crates/ruff_linter/src/rules/airflow/helpers.rs new file mode 100644 index 0000000000..9508153a3d --- /dev/null +++ b/crates/ruff_linter/src/rules/airflow/helpers.rs @@ -0,0 +1,64 @@ +use crate::rules::numpy::helpers::ImportSearcher; +use ruff_python_ast::statement_visitor::StatementVisitor; +use ruff_python_ast::{Expr, ExprName, StmtTry}; +use ruff_python_semantic::Exceptions; +use ruff_python_semantic::SemanticModel; + +#[derive(Debug, Clone, Eq, PartialEq)] +pub(crate) enum Replacement { + None, + Name(&'static str), + Message(&'static str), + AutoImport { + path: &'static str, + name: &'static str, + }, +} + +pub(crate) fn is_guarded_by_try_except( + expr: &Expr, + replacement: &Replacement, + semantic: &SemanticModel, +) -> bool { + match expr { + Expr::Name(ExprName { id, .. }) => { + let Some(binding_id) = semantic.lookup_symbol(id.as_str()) else { + return false; + }; + let binding = semantic.binding(binding_id); + if !binding.is_external() { + return false; + } + if !binding.in_exception_handler() { + return false; + } + let Some(try_node) = binding.source.and_then(|import_id| { + semantic + .statements(import_id) + .find_map(|stmt| stmt.as_try_stmt()) + }) else { + return false; + }; + let suspended_exceptions = Exceptions::from_try_stmt(try_node, semantic); + if !suspended_exceptions + .intersects(Exceptions::IMPORT_ERROR | Exceptions::MODULE_NOT_FOUND_ERROR) + { + return false; + } + try_block_contains_undeprecated_import(try_node, replacement) + } + _ => false, + } +} + +/// Given an [`ast::StmtTry`] node, does the `try` branch of that node +/// contain any [`ast::StmtImportFrom`] nodes that indicate the numpy +/// member is being imported from the non-deprecated location? +fn try_block_contains_undeprecated_import(try_node: &StmtTry, replacement: &Replacement) -> bool { + let Replacement::AutoImport { path, name } = replacement else { + return false; + }; + let mut import_searcher = ImportSearcher::new(path, name); + import_searcher.visit_body(&try_node.body); + import_searcher.found_import +} diff --git a/crates/ruff_linter/src/rules/airflow/mod.rs b/crates/ruff_linter/src/rules/airflow/mod.rs index be4a9fe87d..deec05bae5 100644 --- a/crates/ruff_linter/src/rules/airflow/mod.rs +++ b/crates/ruff_linter/src/rules/airflow/mod.rs @@ -1,4 +1,5 @@ //! Airflow-specific rules. +pub(crate) mod helpers; pub(crate) mod rules; #[cfg(test)] @@ -16,6 +17,7 @@ mod tests { #[test_case(Rule::AirflowDagNoScheduleArgument, Path::new("AIR002.py"))] #[test_case(Rule::Airflow3Removal, Path::new("AIR302_args.py"))] #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))] + #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names_try.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"))] 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 5c203e2aed..3e7a9f0608 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 @@ -1,4 +1,6 @@ use crate::checkers::ast::Checker; +use crate::importer::ImportRequest; +use crate::rules::airflow::helpers::{is_guarded_by_try_except, Replacement}; use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, ViolationMetadata}; use ruff_python_ast::helpers::map_callable; @@ -59,26 +61,22 @@ impl Violation for Airflow3Removal { Replacement::Message(message) => { format!("`{deprecated}` is removed in Airflow 3.0; {message}") } + Replacement::AutoImport { path: _, name: _ } => { + format!("`{deprecated}` is removed in Airflow 3.0") + } } } fn fix_title(&self) -> Option { let Airflow3Removal { replacement, .. } = self; - if let Replacement::Name(name) = replacement { - Some(format!("Use `{name}` instead")) - } else { - None + match replacement { + Replacement::Name(name) => Some(format!("Use `{name}` instead")), + Replacement::AutoImport { path, name } => Some(format!("Use `{path}.{name}` instead")), + _ => None, } } } -#[derive(Debug, Eq, PartialEq)] -enum Replacement { - None, - Name(&'static str), - Message(&'static str), -} - /// AIR302 pub(crate) fn airflow_3_removal_expr(checker: &Checker, expr: &Expr) { if !checker.semantic().seen_module(Modules::AIRFLOW) { @@ -559,7 +557,9 @@ fn check_method(checker: &Checker, call_expr: &ExprCall) { /// SubDagOperator() /// ``` fn check_name(checker: &Checker, expr: &Expr, range: TextRange) { - let Some(qualified_name) = checker.semantic().resolve_qualified_name(expr) else { + let semantic = checker.semantic(); + + let Some(qualified_name) = semantic.resolve_qualified_name(expr) else { return; }; @@ -574,7 +574,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) { Replacement::Name("airflow.api_connexion.security.requires_access_*") } ["airflow", "api_connexion", "security", "requires_access_dataset"] => { - Replacement::Name("airflow.api_connexion.security.requires_access_asset") + Replacement::AutoImport { + path: "airflow.api_connexion.security", + name: "requires_access_asset", + } } // airflow.auth.managers @@ -608,9 +611,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) { } // airflow.datasets - ["airflow", "Dataset"] | ["airflow", "datasets", "Dataset"] => { - Replacement::Name("airflow.sdk.Asset") - } + ["airflow", "Dataset"] | ["airflow", "datasets", "Dataset"] => Replacement::AutoImport { + path: "airflow.sdk", + name: "Asset", + }, ["airflow", "datasets", rest @ ..] => match &rest { ["DatasetAliasEvent"] => Replacement::None, ["DatasetAlias"] => Replacement::Name("airflow.sdk.AssetAlias"), @@ -904,13 +908,31 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) { _ => return, }; - checker.report_diagnostic(Diagnostic::new( + if is_guarded_by_try_except(expr, &replacement, semantic) { + return; + } + + let mut diagnostic = Diagnostic::new( Airflow3Removal { deprecated: qualified_name.to_string(), - replacement, + replacement: replacement.clone(), }, range, - )); + ); + + if let Replacement::AutoImport { path, name } = replacement { + diagnostic.try_set_fix(|| { + let (import_edit, binding) = checker.importer().get_or_import_symbol( + &ImportRequest::import_from(path, name), + expr.start(), + checker.semantic(), + )?; + let replacement_edit = Edit::range_replacement(binding, range); + Ok(Fix::safe_edits(import_edit, [replacement_edit])) + }); + }; + + checker.report_diagnostic(diagnostic); } /// Check whether a customized Airflow plugin contains removed extensions. diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_class_attribute.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_class_attribute.py.snap index 3edb0bf022..18fb9d67e7 100644 --- a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_class_attribute.py.snap +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_class_attribute.py.snap @@ -1,7 +1,7 @@ --- source: crates/ruff_linter/src/rules/airflow/mod.rs --- -AIR302_class_attribute.py:24:21: AIR302 `airflow.Dataset` is removed in Airflow 3.0 +AIR302_class_attribute.py:24:21: AIR302 [*] `airflow.Dataset` is removed in Airflow 3.0 | 23 | # airflow.Dataset 24 | dataset_from_root = DatasetFromRoot() @@ -11,6 +11,19 @@ AIR302_class_attribute.py:24:21: AIR302 `airflow.Dataset` is removed in Airflow | = help: Use `airflow.sdk.Asset` instead +ℹ Safe fix +19 19 | from airflow.providers_manager import ProvidersManager +20 20 | from airflow.secrets.base_secrets import BaseSecretsBackend +21 21 | from airflow.secrets.local_filesystem import LocalFilesystemBackend + 22 |+from airflow.sdk import Asset +22 23 | +23 24 | # airflow.Dataset +24 |-dataset_from_root = DatasetFromRoot() + 25 |+dataset_from_root = Asset() +25 26 | dataset_from_root.iter_datasets() +26 27 | dataset_from_root.iter_dataset_aliases() +27 28 | + AIR302_class_attribute.py:25:19: AIR302 [*] `iter_datasets` is removed in Airflow 3.0 | 23 | # airflow.Dataset @@ -52,7 +65,7 @@ AIR302_class_attribute.py:26:19: AIR302 [*] `iter_dataset_aliases` is removed in 28 28 | # airflow.datasets 29 29 | dataset_to_test_method_call = Dataset() -AIR302_class_attribute.py:29:31: AIR302 `airflow.datasets.Dataset` is removed in Airflow 3.0 +AIR302_class_attribute.py:29:31: AIR302 [*] `airflow.datasets.Dataset` is removed in Airflow 3.0 | 28 | # airflow.datasets 29 | dataset_to_test_method_call = Dataset() @@ -62,6 +75,24 @@ AIR302_class_attribute.py:29:31: AIR302 `airflow.datasets.Dataset` is removed in | = help: Use `airflow.sdk.Asset` instead +ℹ Safe fix +19 19 | from airflow.providers_manager import ProvidersManager +20 20 | from airflow.secrets.base_secrets import BaseSecretsBackend +21 21 | from airflow.secrets.local_filesystem import LocalFilesystemBackend + 22 |+from airflow.sdk import Asset +22 23 | +23 24 | # airflow.Dataset +24 25 | dataset_from_root = DatasetFromRoot() +-------------------------------------------------------------------------------- +26 27 | dataset_from_root.iter_dataset_aliases() +27 28 | +28 29 | # airflow.datasets +29 |-dataset_to_test_method_call = Dataset() + 30 |+dataset_to_test_method_call = Asset() +30 31 | dataset_to_test_method_call.iter_datasets() +31 32 | dataset_to_test_method_call.iter_dataset_aliases() +32 33 | + AIR302_class_attribute.py:30:29: AIR302 [*] `iter_datasets` is removed in Airflow 3.0 | 28 | # airflow.datasets diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap index 27b5341219..9199ee09b0 100644 --- a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap @@ -64,7 +64,7 @@ AIR302_names.py:121:39: AIR302 `airflow.PY312` is removed in Airflow 3.0 | = help: Use `sys.version_info` instead -AIR302_names.py:122:1: AIR302 `airflow.Dataset` is removed in Airflow 3.0 +AIR302_names.py:122:1: AIR302 [*] `airflow.Dataset` is removed in Airflow 3.0 | 120 | # airflow root 121 | PY36, PY37, PY38, PY39, PY310, PY311, PY312 @@ -75,6 +75,20 @@ AIR302_names.py:122:1: AIR302 `airflow.Dataset` is removed in Airflow 3.0 | = help: Use `airflow.sdk.Asset` instead +ℹ Safe fix +116 116 | from airflow.utils.trigger_rule import TriggerRule +117 117 | from airflow.www.auth import has_access, has_access_dataset +118 118 | from airflow.www.utils import get_sensitive_variables_fields, should_hide_value_for_key + 119 |+from airflow.sdk import Asset +119 120 | +120 121 | # airflow root +121 122 | PY36, PY37, PY38, PY39, PY310, PY311, PY312 +122 |-DatasetFromRoot() + 123 |+Asset() +123 124 | +124 125 | # airflow.api_connexion.security +125 126 | requires_access, requires_access_dataset + AIR302_names.py:125:1: AIR302 `airflow.api_connexion.security.requires_access` is removed in Airflow 3.0 | 124 | # airflow.api_connexion.security @@ -85,7 +99,7 @@ AIR302_names.py:125:1: AIR302 `airflow.api_connexion.security.requires_access` i | = help: Use `airflow.api_connexion.security.requires_access_*` instead -AIR302_names.py:125:18: AIR302 `airflow.api_connexion.security.requires_access_dataset` is removed in Airflow 3.0 +AIR302_names.py:125:18: AIR302 [*] `airflow.api_connexion.security.requires_access_dataset` is removed in Airflow 3.0 | 124 | # airflow.api_connexion.security 125 | requires_access, requires_access_dataset @@ -95,6 +109,25 @@ AIR302_names.py:125:18: AIR302 `airflow.api_connexion.security.requires_access_d | = help: Use `airflow.api_connexion.security.requires_access_asset` instead +ℹ Safe fix +12 12 | from airflow import ( +13 13 | Dataset as DatasetFromRoot, +14 14 | ) +15 |-from airflow.api_connexion.security import requires_access, requires_access_dataset + 15 |+from airflow.api_connexion.security import requires_access, requires_access_dataset, requires_access_asset +16 16 | from airflow.auth.managers.base_auth_manager import is_authorized_dataset +17 17 | from airflow.auth.managers.models.resource_details import DatasetDetails +18 18 | from airflow.configuration import ( +-------------------------------------------------------------------------------- +122 122 | DatasetFromRoot() +123 123 | +124 124 | # airflow.api_connexion.security +125 |-requires_access, requires_access_dataset + 125 |+requires_access, requires_access_asset +126 126 | +127 127 | # airflow.auth.managers +128 128 | is_authorized_dataset + AIR302_names.py:128:1: AIR302 `airflow.auth.managers.base_auth_manager.is_authorized_dataset` is removed in Airflow 3.0 | 127 | # airflow.auth.managers @@ -188,7 +221,7 @@ AIR302_names.py:136:1: AIR302 `airflow.contrib.aws_athena_hook.AWSAthenaHook` is 138 | # airflow.datasets | -AIR302_names.py:139:1: AIR302 `airflow.datasets.Dataset` is removed in Airflow 3.0 +AIR302_names.py:139:1: AIR302 [*] `airflow.datasets.Dataset` is removed in Airflow 3.0 | 138 | # airflow.datasets 139 | Dataset() @@ -198,6 +231,24 @@ AIR302_names.py:139:1: AIR302 `airflow.datasets.Dataset` is removed in Airflow 3 | = help: Use `airflow.sdk.Asset` instead +ℹ Safe fix +116 116 | from airflow.utils.trigger_rule import TriggerRule +117 117 | from airflow.www.auth import has_access, has_access_dataset +118 118 | from airflow.www.utils import get_sensitive_variables_fields, should_hide_value_for_key + 119 |+from airflow.sdk import Asset +119 120 | +120 121 | # airflow root +121 122 | PY36, PY37, PY38, PY39, PY310, PY311, PY312 +-------------------------------------------------------------------------------- +136 137 | AWSAthenaHook() +137 138 | +138 139 | # airflow.datasets +139 |-Dataset() + 140 |+Asset() +140 141 | DatasetAlias() +141 142 | DatasetAliasEvent() +142 143 | DatasetAll() + AIR302_names.py:140:1: AIR302 `airflow.datasets.DatasetAlias` is removed in Airflow 3.0 | 138 | # airflow.datasets diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names_try.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names_try.py.snap new file mode 100644 index 0000000000..707c6a2775 --- /dev/null +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names_try.py.snap @@ -0,0 +1,5 @@ +--- +source: crates/ruff_linter/src/rules/airflow/mod.rs +snapshot_kind: text +--- + diff --git a/crates/ruff_linter/src/rules/numpy/helpers.rs b/crates/ruff_linter/src/rules/numpy/helpers.rs new file mode 100644 index 0000000000..39ee1c1767 --- /dev/null +++ b/crates/ruff_linter/src/rules/numpy/helpers.rs @@ -0,0 +1,45 @@ +use ruff_python_ast::statement_visitor::StatementVisitor; +use ruff_python_ast::{statement_visitor, Alias, Stmt, StmtImportFrom}; + +/// AST visitor that searches an AST tree for [`ast::StmtImportFrom`] nodes +/// that match a certain [`QualifiedName`]. +pub(crate) struct ImportSearcher<'a> { + module: &'a str, + name: &'a str, + pub found_import: bool, +} + +impl<'a> ImportSearcher<'a> { + pub(crate) fn new(module: &'a str, name: &'a str) -> Self { + Self { + module, + name, + found_import: false, + } + } +} +impl StatementVisitor<'_> for ImportSearcher<'_> { + fn visit_stmt(&mut self, stmt: &Stmt) { + if self.found_import { + return; + } + if let Stmt::ImportFrom(StmtImportFrom { module, names, .. }) = stmt { + if module.as_ref().is_some_and(|module| module == self.module) + && names.iter().any(|Alias { name, .. }| name == self.name) + { + self.found_import = true; + return; + } + } + statement_visitor::walk_stmt(self, stmt); + } + + fn visit_body(&mut self, body: &[ruff_python_ast::Stmt]) { + for stmt in body { + self.visit_stmt(stmt); + if self.found_import { + return; + } + } + } +} diff --git a/crates/ruff_linter/src/rules/numpy/mod.rs b/crates/ruff_linter/src/rules/numpy/mod.rs index fdf65ba09e..7a313eac2e 100644 --- a/crates/ruff_linter/src/rules/numpy/mod.rs +++ b/crates/ruff_linter/src/rules/numpy/mod.rs @@ -1,4 +1,5 @@ //! NumPy-specific rules. +pub(crate) mod helpers; pub(crate) mod rules; #[cfg(test)] diff --git a/crates/ruff_linter/src/rules/numpy/rules/numpy_2_0_deprecation.rs b/crates/ruff_linter/src/rules/numpy/rules/numpy_2_0_deprecation.rs index 730a6a5137..64ad988938 100644 --- a/crates/ruff_linter/src/rules/numpy/rules/numpy_2_0_deprecation.rs +++ b/crates/ruff_linter/src/rules/numpy/rules/numpy_2_0_deprecation.rs @@ -1,3 +1,4 @@ +use crate::rules::numpy::helpers::ImportSearcher; use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, ViolationMetadata}; use ruff_python_ast::name::{QualifiedName, QualifiedNameBuilder}; @@ -892,49 +893,3 @@ fn try_block_contains_undeprecated_import( import_searcher.visit_body(&try_node.body); import_searcher.found_import } - -/// AST visitor that searches an AST tree for [`ast::StmtImportFrom`] nodes -/// that match a certain [`QualifiedName`]. -struct ImportSearcher<'a> { - module: &'a str, - name: &'a str, - found_import: bool, -} - -impl<'a> ImportSearcher<'a> { - fn new(module: &'a str, name: &'a str) -> Self { - Self { - module, - name, - found_import: false, - } - } -} - -impl StatementVisitor<'_> for ImportSearcher<'_> { - fn visit_stmt(&mut self, stmt: &ast::Stmt) { - if self.found_import { - return; - } - if let ast::Stmt::ImportFrom(ast::StmtImportFrom { module, names, .. }) = stmt { - if module.as_ref().is_some_and(|module| module == self.module) - && names - .iter() - .any(|ast::Alias { name, .. }| name == self.name) - { - self.found_import = true; - return; - } - } - ast::statement_visitor::walk_stmt(self, stmt); - } - - fn visit_body(&mut self, body: &[ruff_python_ast::Stmt]) { - for stmt in body { - self.visit_stmt(stmt); - if self.found_import { - return; - } - } - } -}