Skip to content

Commit 6e9fb9a

Browse files
authored
[airflow] Skip attribute check in try catch block (AIR301) (astral-sh#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`
1 parent a507c1b commit 6e9fb9a

File tree

4 files changed

+114
-55
lines changed

4 files changed

+114
-55
lines changed

crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names_try.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,12 @@
66
from airflow.datasets import Dataset as Asset
77

88
Asset
9+
10+
try:
11+
from airflow.sdk import Asset
12+
except ModuleNotFoundError:
13+
from airflow import datasets
14+
15+
Asset = datasets.Dataset
16+
17+
asset = Asset()

crates/ruff_linter/src/rules/airflow/helpers.rs

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1-
use crate::rules::numpy::helpers::ImportSearcher;
1+
use crate::rules::numpy::helpers::{AttributeSearcher, ImportSearcher};
2+
use ruff_python_ast::name::QualifiedNameBuilder;
23
use ruff_python_ast::statement_visitor::StatementVisitor;
4+
use ruff_python_ast::visitor::Visitor;
35
use ruff_python_ast::{Expr, ExprName, StmtTry};
46
use ruff_python_semantic::Exceptions;
57
use ruff_python_semantic::SemanticModel;
@@ -47,6 +49,22 @@ pub(crate) fn is_guarded_by_try_except(
4749
semantic: &SemanticModel,
4850
) -> bool {
4951
match expr {
52+
Expr::Attribute(_) => {
53+
if !semantic.in_exception_handler() {
54+
return false;
55+
}
56+
let Some(try_node) = semantic
57+
.current_statements()
58+
.find_map(|stmt| stmt.as_try_stmt())
59+
else {
60+
return false;
61+
};
62+
let suspended_exceptions = Exceptions::from_try_stmt(try_node, semantic);
63+
if !suspended_exceptions.contains(Exceptions::ATTRIBUTE_ERROR) {
64+
return false;
65+
}
66+
try_block_contains_undeprecated_attribute(try_node, replacement, semantic)
67+
}
5068
Expr::Name(ExprName { id, .. }) => {
5169
let Some(binding_id) = semantic.lookup_symbol(id.as_str()) else {
5270
return false;
@@ -78,7 +96,31 @@ pub(crate) fn is_guarded_by_try_except(
7896
}
7997

8098
/// Given an [`ast::StmtTry`] node, does the `try` branch of that node
81-
/// contain any [`ast::StmtImportFrom`] nodes that indicate the numpy
99+
/// contain any [`ast::ExprAttribute`] nodes that indicate the airflow
100+
/// member is being accessed from the non-deprecated location?
101+
fn try_block_contains_undeprecated_attribute(
102+
try_node: &StmtTry,
103+
replacement: &Replacement,
104+
semantic: &SemanticModel,
105+
) -> bool {
106+
let Replacement::AutoImport { module, name } = replacement else {
107+
return false;
108+
};
109+
let undeprecated_qualified_name = {
110+
let mut builder = QualifiedNameBuilder::default();
111+
for part in module.split('.') {
112+
builder.push(part);
113+
}
114+
builder.push(name);
115+
builder.build()
116+
};
117+
let mut attribute_searcher = AttributeSearcher::new(undeprecated_qualified_name, semantic);
118+
attribute_searcher.visit_body(&try_node.body);
119+
attribute_searcher.found_attribute
120+
}
121+
122+
/// Given an [`ast::StmtTry`] node, does the `try` branch of that node
123+
/// contain any [`ast::StmtImportFrom`] nodes that indicate the airflow
82124
/// member is being imported from the non-deprecated location?
83125
fn try_block_contains_undeprecated_import(try_node: &StmtTry, replacement: &Replacement) -> bool {
84126
let Replacement::AutoImport { module, name } = replacement else {

crates/ruff_linter/src/rules/numpy/helpers.rs

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
1+
use ruff_python_ast::name::QualifiedName;
12
use ruff_python_ast::statement_visitor::StatementVisitor;
3+
use ruff_python_ast::visitor::Visitor;
4+
use ruff_python_ast::visitor::{walk_expr, walk_stmt};
5+
use ruff_python_ast::Expr;
26
use ruff_python_ast::{statement_visitor, Alias, Stmt, StmtImportFrom};
7+
use ruff_python_semantic::SemanticModel;
38

49
/// AST visitor that searches an AST tree for [`ast::StmtImportFrom`] nodes
510
/// that match a certain [`QualifiedName`].
@@ -43,3 +48,57 @@ impl StatementVisitor<'_> for ImportSearcher<'_> {
4348
}
4449
}
4550
}
51+
52+
/// AST visitor that searches an AST tree for [`ast::ExprAttribute`] nodes
53+
/// that match a certain [`QualifiedName`].
54+
pub(crate) struct AttributeSearcher<'a> {
55+
attribute_to_find: QualifiedName<'a>,
56+
semantic: &'a SemanticModel<'a>,
57+
pub found_attribute: bool,
58+
}
59+
60+
impl<'a> AttributeSearcher<'a> {
61+
pub(crate) fn new(
62+
attribute_to_find: QualifiedName<'a>,
63+
semantic: &'a SemanticModel<'a>,
64+
) -> Self {
65+
Self {
66+
attribute_to_find,
67+
semantic,
68+
found_attribute: false,
69+
}
70+
}
71+
}
72+
73+
impl Visitor<'_> for AttributeSearcher<'_> {
74+
fn visit_expr(&mut self, expr: &'_ Expr) {
75+
if self.found_attribute {
76+
return;
77+
}
78+
if expr.is_attribute_expr()
79+
&& self
80+
.semantic
81+
.resolve_qualified_name(expr)
82+
.is_some_and(|qualified_name| qualified_name == self.attribute_to_find)
83+
{
84+
self.found_attribute = true;
85+
return;
86+
}
87+
walk_expr(self, expr);
88+
}
89+
90+
fn visit_stmt(&mut self, stmt: &ruff_python_ast::Stmt) {
91+
if !self.found_attribute {
92+
walk_stmt(self, stmt);
93+
}
94+
}
95+
96+
fn visit_body(&mut self, body: &[ruff_python_ast::Stmt]) {
97+
for stmt in body {
98+
self.visit_stmt(stmt);
99+
if self.found_attribute {
100+
return;
101+
}
102+
}
103+
}
104+
}

crates/ruff_linter/src/rules/numpy/rules/numpy_2_0_deprecation.rs

Lines changed: 2 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
use crate::rules::numpy::helpers::ImportSearcher;
1+
use crate::rules::numpy::helpers::{AttributeSearcher, ImportSearcher};
22
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
33
use ruff_macros::{derive_message_formats, ViolationMetadata};
4-
use ruff_python_ast::name::{QualifiedName, QualifiedNameBuilder};
4+
use ruff_python_ast::name::QualifiedNameBuilder;
55
use ruff_python_ast::statement_visitor::StatementVisitor;
66
use ruff_python_ast::visitor::Visitor;
77
use ruff_python_ast::{self as ast, Expr};
@@ -823,57 +823,6 @@ fn try_block_contains_undeprecated_attribute(
823823
attribute_searcher.found_attribute
824824
}
825825

826-
/// AST visitor that searches an AST tree for [`ast::ExprAttribute`] nodes
827-
/// that match a certain [`QualifiedName`].
828-
struct AttributeSearcher<'a> {
829-
attribute_to_find: QualifiedName<'a>,
830-
semantic: &'a SemanticModel<'a>,
831-
found_attribute: bool,
832-
}
833-
834-
impl<'a> AttributeSearcher<'a> {
835-
fn new(attribute_to_find: QualifiedName<'a>, semantic: &'a SemanticModel<'a>) -> Self {
836-
Self {
837-
attribute_to_find,
838-
semantic,
839-
found_attribute: false,
840-
}
841-
}
842-
}
843-
844-
impl Visitor<'_> for AttributeSearcher<'_> {
845-
fn visit_expr(&mut self, expr: &'_ Expr) {
846-
if self.found_attribute {
847-
return;
848-
}
849-
if expr.is_attribute_expr()
850-
&& self
851-
.semantic
852-
.resolve_qualified_name(expr)
853-
.is_some_and(|qualified_name| qualified_name == self.attribute_to_find)
854-
{
855-
self.found_attribute = true;
856-
return;
857-
}
858-
ast::visitor::walk_expr(self, expr);
859-
}
860-
861-
fn visit_stmt(&mut self, stmt: &ruff_python_ast::Stmt) {
862-
if !self.found_attribute {
863-
ast::visitor::walk_stmt(self, stmt);
864-
}
865-
}
866-
867-
fn visit_body(&mut self, body: &[ruff_python_ast::Stmt]) {
868-
for stmt in body {
869-
self.visit_stmt(stmt);
870-
if self.found_attribute {
871-
return;
872-
}
873-
}
874-
}
875-
}
876-
877826
/// Given an [`ast::StmtTry`] node, does the `try` branch of that node
878827
/// contain any [`ast::StmtImportFrom`] nodes that indicate the numpy
879828
/// member is being imported from the non-deprecated location?

0 commit comments

Comments
 (0)