[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`
This commit is contained in:
Wei Lee 2025-05-05 22:01:05 +08:00 committed by GitHub
parent a507c1b8b3
commit 6e9fb9af38
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 114 additions and 55 deletions

View file

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

View file

@ -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 {

View file

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

View file

@ -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?