[airflow] Apply try-catch guard to all AIR3 rules (AIR3) (#17887)

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

If a try-catch block guards the names, we don't raise warnings. During
this change, I discovered that some of the replacement types were
missed. Thus, I extend the fix to types other than AutoImport as well

## Test Plan

<!-- How was it tested? -->

Test fixtures are added and updated.

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,13 +937,17 @@ 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 let Some((module, name)) = match &replacement {
+        ProviderReplacement::AutoImport { module, name, .. } => Some((module, *name)),
+        ProviderReplacement::SourceModuleMovedToProvider { module, name, .. } => {
+            Some((module, name.as_str()))
+        }
+        _ => None,
+    } {
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
         diagnostic.try_set_fix(|| {
             let (import_edit, binding) = checker.importer().get_or_import_symbol(
                 &ImportRequest::import_from(module, name),
@@ -954,6 +958,5 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
             Ok(Fix::safe_edits(import_edit, [replacement_edit]))
         });
     }
-
     checker.report_diagnostic(diagnostic);
 }

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

@@ -865,19 +865,22 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         _ => return,
     };
 
-    if is_guarded_by_try_except(expr, &replacement, semantic) {
-        return;
-    }
-
     let mut diagnostic = Diagnostic::new(
         Airflow3Removal {
             deprecated: qualified_name.to_string(),
             replacement: replacement.clone(),
         },
         range,
     );
-
-    if let Replacement::AutoImport { module, name } = replacement {
+    let semantic = checker.semantic();
+    if let Some((module, name)) = match &replacement {
+        Replacement::AutoImport { module, name } => Some((module, *name)),
+        Replacement::SourceModuleMoved { module, name } => Some((module, name.as_str())),
+        _ => None,
+    } {
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
         diagnostic.try_set_fix(|| {
             let (import_edit, binding) = checker.importer().get_or_import_symbol(
                 &ImportRequest::import_from(module, name),

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};
@@ -279,13 +279,17 @@ 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 let Some((module, name)) = match &replacement {
+        ProviderReplacement::AutoImport { module, name, .. } => Some((module, *name)),
+        ProviderReplacement::SourceModuleMovedToProvider { module, name, .. } => {
+            Some((module, name.as_str()))
+        }
+        _ => None,
+    } {
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
         diagnostic.try_set_fix(|| {
             let (import_edit, binding) = checker.importer().get_or_import_symbol(
                 &ImportRequest::import_from(module, name),

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

@@ -283,10 +283,6 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         _ => return,
     };
 
-    if is_guarded_by_try_except(expr, &replacement, semantic) {
-        return;
-    }
-
     let mut diagnostic = Diagnostic::new(
         Airflow3SuggestedUpdate {
             deprecated: qualified_name.to_string(),
@@ -295,7 +291,15 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         range,
     );
 
-    if let Replacement::AutoImport { module, name } = replacement {
+    let semantic = checker.semantic();
+    if let Some((module, name)) = match &replacement {
+        Replacement::AutoImport { module, name } => Some((module, *name)),
+        Replacement::SourceModuleMoved { module, name } => Some((module, name.as_str())),
+        _ => None,
+    } {
+        if is_guarded_by_try_except(expr, module, name, semantic) {
+            return;
+        }
         diagnostic.try_set_fix(|| {
             let (import_edit, binding) = checker.importer().get_or_import_symbol(
                 &ImportRequest::import_from(module, name),

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_try.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/airflow/mod.rs
+---
+

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR311_AIR311_names.py.snap

@@ -258,7 +258,7 @@ AIR311_names.py:49:1: AIR311 `airflow.models.Connection` is removed in Airflow 3
    |
    = help: Use `airflow.sdk.Connection` instead
 
-AIR311_names.py:50:1: AIR311 `airflow.models.DAG` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:50:1: AIR311 [*] `airflow.models.DAG` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 48 | # airflow.models
 49 | Connection()
@@ -268,6 +268,24 @@ AIR311_names.py:50:1: AIR311 `airflow.models.DAG` is removed in Airflow 3.0; It
    |
    = help: Use `airflow.sdk.DAG` instead
 
+ℹ Safe fix
+22 22 | from airflow.models.dag import DAG as DAGFromDag
+23 23 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+24 24 | from airflow.utils.dag_parsing_context import get_parsing_context
+   25 |+from airflow.sdk import DAG
+25 26 | 
+26 27 | # airflow
+27 28 | DatasetFromRoot()
+--------------------------------------------------------------------------------
+47 48 | 
+48 49 | # airflow.models
+49 50 | Connection()
+50    |-DAGFromModel()
+   51 |+DAG()
+51 52 | Variable()
+52 53 | 
+53 54 | # airflow.models.baseoperator
+
 AIR311_names.py:51:1: AIR311 `airflow.models.Variable` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 49 | Connection()
@@ -310,7 +328,7 @@ AIR311_names.py:56:1: AIR311 `airflow.models.baseoperator.cross_downstream` is r
    |
    = help: Use `airflow.sdk.cross_downstream` instead
 
-AIR311_names.py:62:1: AIR311 `airflow.models.dag.DAG` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:62:1: AIR311 [*] `airflow.models.dag.DAG` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 61 | # airflow.models.dag
 62 | DAGFromDag()
@@ -320,6 +338,24 @@ AIR311_names.py:62:1: AIR311 `airflow.models.dag.DAG` is removed in Airflow 3.0;
    |
    = help: Use `airflow.sdk.DAG` instead
 
+ℹ Safe fix
+22 22 | from airflow.models.dag import DAG as DAGFromDag
+23 23 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+24 24 | from airflow.utils.dag_parsing_context import get_parsing_context
+   25 |+from airflow.sdk import DAG
+25 26 | 
+26 27 | # airflow
+27 28 | DatasetFromRoot()
+--------------------------------------------------------------------------------
+59 60 | BaseOperatorLink()
+60 61 | 
+61 62 | # airflow.models.dag
+62    |-DAGFromDag()
+   63 |+DAG()
+63 64 | # airflow.timetables.datasets
+64 65 | DatasetOrTimeSchedule()
+65 66 | 
+
 AIR311_names.py:64:1: AIR311 [*] `airflow.timetables.datasets.DatasetOrTimeSchedule` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 62 | DAGFromDag()

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR311_AIR311_try.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/airflow/mod.rs
+---
+

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR312_AIR312_try.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/airflow/mod.rs
+---
+