From 6e9fb9af38fbc148dde21eff3c89c80eef43d330 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Mon, 5 May 2025 22:01:05 +0800 Subject: [PATCH] [`airflow`] Skip attribute check in try catch block (`AIR301`) (#17790) ## Summary Skip attribute check in try catch block (`AIR301`) ## Test Plan update `crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names_try.py` --- .../test/fixtures/airflow/AIR301_names_try.py | 9 +++ .../ruff_linter/src/rules/airflow/helpers.rs | 46 ++++++++++++++- crates/ruff_linter/src/rules/numpy/helpers.rs | 59 +++++++++++++++++++ .../numpy/rules/numpy_2_0_deprecation.rs | 55 +---------------- 4 files changed, 114 insertions(+), 55 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names_try.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names_try.py index 5d4c3da974..113fe88831 100644 --- a/crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names_try.py +++ b/crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names_try.py @@ -6,3 +6,12 @@ except ModuleNotFoundError: from airflow.datasets import Dataset as Asset Asset + +try: + from airflow.sdk import Asset +except ModuleNotFoundError: + from airflow import datasets + + Asset = datasets.Dataset + +asset = Asset() diff --git a/crates/ruff_linter/src/rules/airflow/helpers.rs b/crates/ruff_linter/src/rules/airflow/helpers.rs index 9cd5818b8f..c5468943be 100644 --- a/crates/ruff_linter/src/rules/airflow/helpers.rs +++ b/crates/ruff_linter/src/rules/airflow/helpers.rs @@ -1,5 +1,7 @@ -use crate::rules::numpy::helpers::ImportSearcher; +use crate::rules::numpy::helpers::{AttributeSearcher, ImportSearcher}; +use ruff_python_ast::name::QualifiedNameBuilder; use ruff_python_ast::statement_visitor::StatementVisitor; +use ruff_python_ast::visitor::Visitor; use ruff_python_ast::{Expr, ExprName, StmtTry}; use ruff_python_semantic::Exceptions; use ruff_python_semantic::SemanticModel; @@ -47,6 +49,22 @@ pub(crate) fn is_guarded_by_try_except( semantic: &SemanticModel, ) -> bool { match expr { + Expr::Attribute(_) => { + if !semantic.in_exception_handler() { + return false; + } + let Some(try_node) = semantic + .current_statements() + .find_map(|stmt| stmt.as_try_stmt()) + else { + return false; + }; + let suspended_exceptions = Exceptions::from_try_stmt(try_node, semantic); + if !suspended_exceptions.contains(Exceptions::ATTRIBUTE_ERROR) { + return false; + } + try_block_contains_undeprecated_attribute(try_node, replacement, semantic) + } Expr::Name(ExprName { id, .. }) => { let Some(binding_id) = semantic.lookup_symbol(id.as_str()) else { return false; @@ -78,7 +96,31 @@ pub(crate) fn is_guarded_by_try_except( } /// Given an [`ast::StmtTry`] node, does the `try` branch of that node -/// contain any [`ast::StmtImportFrom`] nodes that indicate the numpy +/// contain any [`ast::ExprAttribute`] nodes that indicate the airflow +/// member is being accessed from the non-deprecated location? +fn try_block_contains_undeprecated_attribute( + try_node: &StmtTry, + replacement: &Replacement, + semantic: &SemanticModel, +) -> bool { + let Replacement::AutoImport { module, name } = replacement else { + return false; + }; + let undeprecated_qualified_name = { + let mut builder = QualifiedNameBuilder::default(); + for part in module.split('.') { + builder.push(part); + } + builder.push(name); + builder.build() + }; + let mut attribute_searcher = AttributeSearcher::new(undeprecated_qualified_name, semantic); + attribute_searcher.visit_body(&try_node.body); + attribute_searcher.found_attribute +} + +/// Given an [`ast::StmtTry`] node, does the `try` branch of that node +/// contain any [`ast::StmtImportFrom`] nodes that indicate the airflow /// member is being imported from the non-deprecated location? fn try_block_contains_undeprecated_import(try_node: &StmtTry, replacement: &Replacement) -> bool { let Replacement::AutoImport { module, name } = replacement else { diff --git a/crates/ruff_linter/src/rules/numpy/helpers.rs b/crates/ruff_linter/src/rules/numpy/helpers.rs index 39ee1c1767..e5cef62053 100644 --- a/crates/ruff_linter/src/rules/numpy/helpers.rs +++ b/crates/ruff_linter/src/rules/numpy/helpers.rs @@ -1,5 +1,10 @@ +use ruff_python_ast::name::QualifiedName; use ruff_python_ast::statement_visitor::StatementVisitor; +use ruff_python_ast::visitor::Visitor; +use ruff_python_ast::visitor::{walk_expr, walk_stmt}; +use ruff_python_ast::Expr; use ruff_python_ast::{statement_visitor, Alias, Stmt, StmtImportFrom}; +use ruff_python_semantic::SemanticModel; /// AST visitor that searches an AST tree for [`ast::StmtImportFrom`] nodes /// that match a certain [`QualifiedName`]. @@ -43,3 +48,57 @@ impl StatementVisitor<'_> for ImportSearcher<'_> { } } } + +/// AST visitor that searches an AST tree for [`ast::ExprAttribute`] nodes +/// that match a certain [`QualifiedName`]. +pub(crate) struct AttributeSearcher<'a> { + attribute_to_find: QualifiedName<'a>, + semantic: &'a SemanticModel<'a>, + pub found_attribute: bool, +} + +impl<'a> AttributeSearcher<'a> { + pub(crate) fn new( + attribute_to_find: QualifiedName<'a>, + semantic: &'a SemanticModel<'a>, + ) -> Self { + Self { + attribute_to_find, + semantic, + found_attribute: false, + } + } +} + +impl Visitor<'_> for AttributeSearcher<'_> { + fn visit_expr(&mut self, expr: &'_ Expr) { + if self.found_attribute { + return; + } + if expr.is_attribute_expr() + && self + .semantic + .resolve_qualified_name(expr) + .is_some_and(|qualified_name| qualified_name == self.attribute_to_find) + { + self.found_attribute = true; + return; + } + walk_expr(self, expr); + } + + fn visit_stmt(&mut self, stmt: &ruff_python_ast::Stmt) { + if !self.found_attribute { + 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_attribute { + return; + } + } + } +} 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 5fca95a4f7..38f69c4673 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,7 +1,7 @@ -use crate::rules::numpy::helpers::ImportSearcher; +use crate::rules::numpy::helpers::{AttributeSearcher, ImportSearcher}; use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, ViolationMetadata}; -use ruff_python_ast::name::{QualifiedName, QualifiedNameBuilder}; +use ruff_python_ast::name::QualifiedNameBuilder; use ruff_python_ast::statement_visitor::StatementVisitor; use ruff_python_ast::visitor::Visitor; use ruff_python_ast::{self as ast, Expr}; @@ -823,57 +823,6 @@ fn try_block_contains_undeprecated_attribute( attribute_searcher.found_attribute } -/// AST visitor that searches an AST tree for [`ast::ExprAttribute`] nodes -/// that match a certain [`QualifiedName`]. -struct AttributeSearcher<'a> { - attribute_to_find: QualifiedName<'a>, - semantic: &'a SemanticModel<'a>, - found_attribute: bool, -} - -impl<'a> AttributeSearcher<'a> { - fn new(attribute_to_find: QualifiedName<'a>, semantic: &'a SemanticModel<'a>) -> Self { - Self { - attribute_to_find, - semantic, - found_attribute: false, - } - } -} - -impl Visitor<'_> for AttributeSearcher<'_> { - fn visit_expr(&mut self, expr: &'_ Expr) { - if self.found_attribute { - return; - } - if expr.is_attribute_expr() - && self - .semantic - .resolve_qualified_name(expr) - .is_some_and(|qualified_name| qualified_name == self.attribute_to_find) - { - self.found_attribute = true; - return; - } - ast::visitor::walk_expr(self, expr); - } - - fn visit_stmt(&mut self, stmt: &ruff_python_ast::Stmt) { - if !self.found_attribute { - ast::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_attribute { - return; - } - } - } -} - /// 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?