Skip to content

Commit 58fd21c

Browse files
Lee-WLee-W
committed
feat(AIR3): apply try catch guard to all AIR3 rules
1 parent f82b728 commit 58fd21c

14 files changed

Lines changed: 246 additions & 88 deletions

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

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,8 @@
11
from __future__ import annotations
22

33
try:
4-
from airflow.sdk import Asset
4+
from airflow.assets.manager import AssetManager
55
except ModuleNotFoundError:
6-
from airflow.datasets import Dataset as Asset
6+
from airflow.datasets.manager import DatasetManager as AssetManager
77

8-
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()
8+
AssetManager()

‎crates/ruff_linter/resources/test/fixtures/airflow/AIR302_try.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
from __future__ import annotations
2+
3+
try:
4+
from airflow.providers.http.operators.http import HttpOperator
5+
except ModuleNotFoundError:
6+
from airflow.operators.http_operator import SimpleHttpOperator as HttpOperator
7+
8+
HttpOperator()

‎crates/ruff_linter/resources/test/fixtures/airflow/AIR311_try.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
from __future__ import annotations
2+
3+
try:
4+
from airflow.sdk import Asset
5+
except ModuleNotFoundError:
6+
from airflow.datasets import Dataset as Asset
7+
8+
Asset()

‎crates/ruff_linter/resources/test/fixtures/airflow/AIR312_try.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
from __future__ import annotations
2+
3+
try:
4+
from airflow.providers.standard.triggers.file import FileTrigger
5+
except ModuleNotFoundError:
6+
from airflow.triggers.file import FileTrigger
7+
8+
FileTrigger()

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

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ pub(crate) enum ProviderReplacement {
4545

4646
pub(crate) fn is_guarded_by_try_except(
4747
expr: &Expr,
48-
replacement: &Replacement,
48+
module: &str,
49+
name: &str,
4950
semantic: &SemanticModel,
5051
) -> bool {
5152
match expr {
@@ -63,7 +64,7 @@ pub(crate) fn is_guarded_by_try_except(
6364
if !suspended_exceptions.contains(Exceptions::ATTRIBUTE_ERROR) {
6465
return false;
6566
}
66-
try_block_contains_undeprecated_attribute(try_node, replacement, semantic)
67+
try_block_contains_undeprecated_attribute(try_node, module, name, semantic)
6768
}
6869
Expr::Name(ExprName { id, .. }) => {
6970
let Some(binding_id) = semantic.lookup_symbol(id.as_str()) else {
@@ -89,7 +90,7 @@ pub(crate) fn is_guarded_by_try_except(
8990
{
9091
return false;
9192
}
92-
try_block_contains_undeprecated_import(try_node, replacement)
93+
try_block_contains_undeprecated_import(try_node, module, name)
9394
}
9495
_ => false,
9596
}
@@ -100,12 +101,10 @@ pub(crate) fn is_guarded_by_try_except(
100101
/// member is being accessed from the non-deprecated location?
101102
fn try_block_contains_undeprecated_attribute(
102103
try_node: &StmtTry,
103-
replacement: &Replacement,
104+
module: &str,
105+
name: &str,
104106
semantic: &SemanticModel,
105107
) -> bool {
106-
let Replacement::AutoImport { module, name } = replacement else {
107-
return false;
108-
};
109108
let undeprecated_qualified_name = {
110109
let mut builder = QualifiedNameBuilder::default();
111110
for part in module.split('.') {
@@ -122,10 +121,7 @@ fn try_block_contains_undeprecated_attribute(
122121
/// Given an [`ast::StmtTry`] node, does the `try` branch of that node
123122
/// contain any [`ast::StmtImportFrom`] nodes that indicate the airflow
124123
/// member is being imported from the non-deprecated location?
125-
fn try_block_contains_undeprecated_import(try_node: &StmtTry, replacement: &Replacement) -> bool {
126-
let Replacement::AutoImport { module, name } = replacement else {
127-
return false;
128-
};
124+
fn try_block_contains_undeprecated_import(try_node: &StmtTry, module: &str, name: &str) -> bool {
129125
let mut import_searcher = ImportSearcher::new(module, name);
130126
import_searcher.visit_body(&try_node.body);
131127
import_searcher.found_import

‎crates/ruff_linter/src/rules/airflow/mod.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,12 @@ mod tests {
4646
#[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_sqlite.py"))]
4747
#[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_zendesk.py"))]
4848
#[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_standard.py"))]
49+
#[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_try.py"))]
4950
#[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_args.py"))]
5051
#[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_names.py"))]
52+
#[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_try.py"))]
5153
#[test_case(Rule::Airflow3SuggestedToMoveToProvider, Path::new("AIR312.py"))]
54+
#[test_case(Rule::Airflow3SuggestedToMoveToProvider, Path::new("AIR312_try.py"))]
5255
fn rules(rule_code: Rule, path: &Path) -> Result<()> {
5356
let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
5457
let diagnostics = test_path(

‎crates/ruff_linter/src/rules/airflow/rules/moved_to_provider_in_3.rs‎

Lines changed: 42 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use crate::importer::ImportRequest;
2-
use crate::rules::airflow::helpers::ProviderReplacement;
2+
use crate::rules::airflow::helpers::{is_guarded_by_try_except, ProviderReplacement};
33
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
44
use ruff_macros::{derive_message_formats, ViolationMetadata};
55
use ruff_python_ast::{Expr, ExprAttribute};
@@ -937,22 +937,47 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
937937
ranged.range(),
938938
);
939939

940-
if let ProviderReplacement::AutoImport {
941-
module,
942-
name,
943-
provider: _,
944-
version: _,
945-
} = replacement
946-
{
947-
diagnostic.try_set_fix(|| {
948-
let (import_edit, binding) = checker.importer().get_or_import_symbol(
949-
&ImportRequest::import_from(module, name),
950-
expr.start(),
951-
checker.semantic(),
952-
)?;
953-
let replacement_edit = Edit::range_replacement(binding, ranged.range());
954-
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
955-
});
940+
let semantic = checker.semantic();
941+
match replacement {
942+
ProviderReplacement::AutoImport {
943+
module,
944+
name,
945+
provider: _,
946+
version: _,
947+
} => {
948+
if is_guarded_by_try_except(expr, module, name, semantic) {
949+
return;
950+
}
951+
diagnostic.try_set_fix(|| {
952+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
953+
&ImportRequest::import_from(module, name),
954+
expr.start(),
955+
checker.semantic(),
956+
)?;
957+
let replacement_edit = Edit::range_replacement(binding, ranged.range());
958+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
959+
});
960+
}
961+
ProviderReplacement::SourceModuleMovedToProvider {
962+
module,
963+
name,
964+
provider: _,
965+
version: _,
966+
} => {
967+
if is_guarded_by_try_except(expr, module, name.as_str(), semantic) {
968+
return;
969+
}
970+
diagnostic.try_set_fix(|| {
971+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
972+
&ImportRequest::import_from(module, name.as_str()),
973+
expr.start(),
974+
checker.semantic(),
975+
)?;
976+
let replacement_edit = Edit::range_replacement(binding, ranged.range());
977+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
978+
});
979+
}
980+
_ => {}
956981
}
957982

958983
checker.report_diagnostic(diagnostic);

‎crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs‎

Lines changed: 31 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -865,28 +865,44 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
865865
_ => return,
866866
};
867867

868-
if is_guarded_by_try_except(expr, &replacement, semantic) {
869-
return;
870-
}
871-
872868
let mut diagnostic = Diagnostic::new(
873869
Airflow3Removal {
874870
deprecated: qualified_name.to_string(),
875871
replacement: replacement.clone(),
876872
},
877873
range,
878874
);
879-
880-
if let Replacement::AutoImport { module, name } = replacement {
881-
diagnostic.try_set_fix(|| {
882-
let (import_edit, binding) = checker.importer().get_or_import_symbol(
883-
&ImportRequest::import_from(module, name),
884-
expr.start(),
885-
checker.semantic(),
886-
)?;
887-
let replacement_edit = Edit::range_replacement(binding, range);
888-
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
889-
});
875+
let semantic = checker.semantic();
876+
match replacement {
877+
Replacement::AutoImport { module, name } => {
878+
if is_guarded_by_try_except(expr, module, name, semantic) {
879+
return;
880+
}
881+
diagnostic.try_set_fix(|| {
882+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
883+
&ImportRequest::import_from(module, name),
884+
expr.start(),
885+
checker.semantic(),
886+
)?;
887+
let replacement_edit = Edit::range_replacement(binding, range);
888+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
889+
});
890+
}
891+
Replacement::SourceModuleMoved { module, name } => {
892+
if is_guarded_by_try_except(expr, module, name.as_str(), semantic) {
893+
return;
894+
}
895+
diagnostic.try_set_fix(|| {
896+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
897+
&ImportRequest::import_from(module, name.as_str()),
898+
expr.start(),
899+
checker.semantic(),
900+
)?;
901+
let replacement_edit = Edit::range_replacement(binding, range);
902+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
903+
});
904+
}
905+
_ => {}
890906
}
891907

892908
checker.report_diagnostic(diagnostic);

‎crates/ruff_linter/src/rules/airflow/rules/suggested_to_move_to_provider_in_3.rs‎

Lines changed: 55 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::importer::ImportRequest;
22

3-
use crate::rules::airflow::helpers::ProviderReplacement;
3+
use crate::rules::airflow::helpers::{is_guarded_by_try_except, ProviderReplacement};
44
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
55
use ruff_macros::{derive_message_formats, ViolationMetadata};
66
use ruff_python_ast::{Expr, ExprAttribute};
@@ -271,6 +271,19 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
271271
_ => return,
272272
};
273273

274+
if let ProviderReplacement::AutoImport {
275+
module,
276+
name,
277+
provider: _,
278+
version: _,
279+
} = replacement
280+
{
281+
let semantic = checker.semantic();
282+
if is_guarded_by_try_except(expr, module, name, semantic) {
283+
return;
284+
}
285+
}
286+
274287
let mut diagnostic = Diagnostic::new(
275288
Airflow3SuggestedToMoveToProvider {
276289
deprecated: qualified_name.to_string(),
@@ -279,22 +292,47 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
279292
ranged.range(),
280293
);
281294

282-
if let ProviderReplacement::AutoImport {
283-
module,
284-
name,
285-
provider: _,
286-
version: _,
287-
} = replacement
288-
{
289-
diagnostic.try_set_fix(|| {
290-
let (import_edit, binding) = checker.importer().get_or_import_symbol(
291-
&ImportRequest::import_from(module, name),
292-
expr.start(),
293-
checker.semantic(),
294-
)?;
295-
let replacement_edit = Edit::range_replacement(binding, ranged.range());
296-
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
297-
});
295+
let semantic = checker.semantic();
296+
match replacement {
297+
ProviderReplacement::AutoImport {
298+
module,
299+
name,
300+
provider: _,
301+
version: _,
302+
} => {
303+
if is_guarded_by_try_except(expr, module, name, semantic) {
304+
return;
305+
}
306+
diagnostic.try_set_fix(|| {
307+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
308+
&ImportRequest::import_from(module, name),
309+
expr.start(),
310+
checker.semantic(),
311+
)?;
312+
let replacement_edit = Edit::range_replacement(binding, ranged.range());
313+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
314+
});
315+
}
316+
ProviderReplacement::SourceModuleMovedToProvider {
317+
module,
318+
name,
319+
provider: _,
320+
version: _,
321+
} => {
322+
if is_guarded_by_try_except(expr, module, name.as_str(), semantic) {
323+
return;
324+
}
325+
diagnostic.try_set_fix(|| {
326+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
327+
&ImportRequest::import_from(module, name.as_str()),
328+
expr.start(),
329+
checker.semantic(),
330+
)?;
331+
let replacement_edit = Edit::range_replacement(binding, ranged.range());
332+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
333+
});
334+
}
335+
_ => {}
298336
}
299337

300338
checker.report_diagnostic(diagnostic);

‎crates/ruff_linter/src/rules/airflow/rules/suggested_to_update_3_0.rs‎

Lines changed: 31 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -275,10 +275,6 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
275275
_ => return,
276276
};
277277

278-
if is_guarded_by_try_except(expr, &replacement, semantic) {
279-
return;
280-
}
281-
282278
let mut diagnostic = Diagnostic::new(
283279
Airflow3SuggestedUpdate {
284280
deprecated: qualified_name.to_string(),
@@ -287,16 +283,37 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
287283
range,
288284
);
289285

290-
if let Replacement::AutoImport { module, name } = replacement {
291-
diagnostic.try_set_fix(|| {
292-
let (import_edit, binding) = checker.importer().get_or_import_symbol(
293-
&ImportRequest::import_from(module, name),
294-
expr.start(),
295-
checker.semantic(),
296-
)?;
297-
let replacement_edit = Edit::range_replacement(binding, range);
298-
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
299-
});
286+
let semantic = checker.semantic();
287+
match replacement {
288+
Replacement::AutoImport { module, name } => {
289+
if is_guarded_by_try_except(expr, module, name, semantic) {
290+
return;
291+
}
292+
diagnostic.try_set_fix(|| {
293+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
294+
&ImportRequest::import_from(module, name),
295+
expr.start(),
296+
checker.semantic(),
297+
)?;
298+
let replacement_edit = Edit::range_replacement(binding, range);
299+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
300+
});
301+
}
302+
Replacement::SourceModuleMoved { module, name } => {
303+
if is_guarded_by_try_except(expr, module, name.as_str(), semantic) {
304+
return;
305+
}
306+
diagnostic.try_set_fix(|| {
307+
let (import_edit, binding) = checker.importer().get_or_import_symbol(
308+
&ImportRequest::import_from(module, name.as_str()),
309+
expr.start(),
310+
checker.semantic(),
311+
)?;
312+
let replacement_edit = Edit::range_replacement(binding, range);
313+
Ok(Fix::safe_edits(import_edit, [replacement_edit]))
314+
});
315+
}
316+
_ => {}
300317
}
301318

302319
checker.report_diagnostic(diagnostic);

0 commit comments

Comments
 (0)