[airflow] Skip attribute check in try catch block (AIR301) (#17790)

<!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->

Skip attribute check in try catch block (`AIR301`)

## Test Plan

<!-- How was it tested? -->
update
`crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names_try.py`

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()

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 {

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;
+            }
+        }
+    }
+}

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?