feat(AIR3): apply try catch guard to all AIR3 rules

Author: Lee-W

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

@@ -1,17 +1,8 @@
 from __future__ import annotations
 
 try:
-    from airflow.sdk import Asset
+    from airflow.assets.manager import AssetManager
 except ModuleNotFoundError:
-    from airflow.datasets import Dataset as Asset
+    from airflow.datasets.manager import DatasetManager as AssetManager
 
-Asset
-
-try:
-    from airflow.sdk import Asset
-except ModuleNotFoundError:
-    from airflow import datasets
-
-    Asset = datasets.Dataset
-
-asset = Asset()
+AssetManager()

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

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+try:
+    from airflow.providers.http.operators.http import HttpOperator
+except ModuleNotFoundError:
+    from airflow.operators.http_operator import SimpleHttpOperator as HttpOperator
+
+HttpOperator()

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

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+try:
+    from airflow.sdk import Asset
+except ModuleNotFoundError:
+    from airflow.datasets import Dataset as Asset
+
+Asset()

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

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+try:
+    from airflow.providers.standard.triggers.file import FileTrigger
+except ModuleNotFoundError:
+    from airflow.triggers.file import FileTrigger
+
+FileTrigger()

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

@@ -45,7 +45,8 @@ pub(crate) enum ProviderReplacement {
 
 pub(crate) fn is_guarded_by_try_except(
     expr: &Expr,
-    replacement: &Replacement,
+    module: &str,
+    name: &str,
     semantic: &SemanticModel,
 ) -> bool {
     match expr {
@@ -63,7 +64,7 @@ pub(crate) fn is_guarded_by_try_except(
             if !suspended_exceptions.contains(Exceptions::ATTRIBUTE_ERROR) {
                 return false;
             }
-            try_block_contains_undeprecated_attribute(try_node, replacement, semantic)
+            try_block_contains_undeprecated_attribute(try_node, module, name, semantic)
         }
         Expr::Name(ExprName { id, .. }) => {
             let Some(binding_id) = semantic.lookup_symbol(id.as_str()) else {
@@ -89,7 +90,7 @@ pub(crate) fn is_guarded_by_try_except(
             {
                 return false;
             }
-            try_block_contains_undeprecated_import(try_node, replacement)
+            try_block_contains_undeprecated_import(try_node, module, name)
         }
         _ => false,
     }
@@ -100,12 +101,10 @@ pub(crate) fn is_guarded_by_try_except(
 /// member is being accessed from the non-deprecated location?
 fn try_block_contains_undeprecated_attribute(
     try_node: &StmtTry,
-    replacement: &Replacement,
+    module: &str,
+    name: &str,
     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('.') {
@@ -122,10 +121,7 @@ fn try_block_contains_undeprecated_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 {
-        return false;
-    };
+fn try_block_contains_undeprecated_import(try_node: &StmtTry, module: &str, name: &str) -> bool {
     let mut import_searcher = ImportSearcher::new(module, name);
     import_searcher.visit_body(&try_node.body);
     import_searcher.found_import

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

@@ -46,9 +46,12 @@ mod tests {
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_sqlite.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_zendesk.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_standard.py"))]
+    #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_try.py"))]
     #[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_args.py"))]
     #[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_names.py"))]
+    #[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_try.py"))]
     #[test_case(Rule::Airflow3SuggestedToMoveToProvider, Path::new("AIR312.py"))]
+    #[test_case(Rule::Airflow3SuggestedToMoveToProvider, Path::new("AIR312_try.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
         let diagnostics = test_path(

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

@@ -1,5 +1,5 @@
 use crate::importer::ImportRequest;
-use crate::rules::airflow::helpers::ProviderReplacement;
+use crate::rules::airflow::helpers::{is_guarded_by_try_except, ProviderReplacement};
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::{Expr, ExprAttribute};
@@ -937,6 +937,19 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
         ranged.range(),
     );
 
+    if let ProviderReplacement::AutoImport {
+        module,
+        name,
+        provider: _,
+        version: _,
+    } = replacement
+    {
+        let semantic = checker.semantic();
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
+    }
+
     if let ProviderReplacement::AutoImport {
         module,
         name,

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

@@ -865,8 +865,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         _ => return,
     };
 
-    if is_guarded_by_try_except(expr, &replacement, semantic) {
-        return;
+    if let Replacement::AutoImport { module, name } = replacement {
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
     }
 
     let mut diagnostic = Diagnostic::new(

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

@@ -1,6 +1,6 @@
 use crate::importer::ImportRequest;
 
-use crate::rules::airflow::helpers::ProviderReplacement;
+use crate::rules::airflow::helpers::{is_guarded_by_try_except, ProviderReplacement};
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::{Expr, ExprAttribute};
@@ -271,6 +271,19 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
         _ => return,
     };
 
+    if let ProviderReplacement::AutoImport {
+        module,
+        name,
+        provider: _,
+        version: _,
+    } = replacement
+    {
+        let semantic = checker.semantic();
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
+    }
+
     let mut diagnostic = Diagnostic::new(
         Airflow3SuggestedToMoveToProvider {
             deprecated: qualified_name.to_string(),

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

@@ -275,8 +275,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         _ => return,
     };
 
-    if is_guarded_by_try_except(expr, &replacement, semantic) {
-        return;
+    if let Replacement::AutoImport { module, name } = replacement {
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
     }
 
     let mut diagnostic = Diagnostic::new(