[airflow] Add autofix infrastructure to AIR302 name checks (#16965)

## Summary

Add autofix infrastructure to `AIR302` name checks and use this logic to 
fix`"airflow", "api_connexion", "security", "requires_access_dataset"`, `"airflow", "Dataset"` and `"airflow",
"datasets", "Dataset"`

## Test Plan

The existing test fixture reflects the update

crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names_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/src/rules/airflow/helpers.rs

@@ -0,0 +1,64 @@
+use crate::rules::numpy::helpers::ImportSearcher;
+use ruff_python_ast::statement_visitor::StatementVisitor;
+use ruff_python_ast::{Expr, ExprName, StmtTry};
+use ruff_python_semantic::Exceptions;
+use ruff_python_semantic::SemanticModel;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub(crate) enum Replacement {
+    None,
+    Name(&'static str),
+    Message(&'static str),
+    AutoImport {
+        path: &'static str,
+        name: &'static str,
+    },
+}
+
+pub(crate) fn is_guarded_by_try_except(
+    expr: &Expr,
+    replacement: &Replacement,
+    semantic: &SemanticModel,
+) -> bool {
+    match expr {
+        Expr::Name(ExprName { id, .. }) => {
+            let Some(binding_id) = semantic.lookup_symbol(id.as_str()) else {
+                return false;
+            };
+            let binding = semantic.binding(binding_id);
+            if !binding.is_external() {
+                return false;
+            }
+            if !binding.in_exception_handler() {
+                return false;
+            }
+            let Some(try_node) = binding.source.and_then(|import_id| {
+                semantic
+                    .statements(import_id)
+                    .find_map(|stmt| stmt.as_try_stmt())
+            }) else {
+                return false;
+            };
+            let suspended_exceptions = Exceptions::from_try_stmt(try_node, semantic);
+            if !suspended_exceptions
+                .intersects(Exceptions::IMPORT_ERROR | Exceptions::MODULE_NOT_FOUND_ERROR)
+            {
+                return false;
+            }
+            try_block_contains_undeprecated_import(try_node, replacement)
+        }
+        _ => false,
+    }
+}
+
+/// 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?
+fn try_block_contains_undeprecated_import(try_node: &StmtTry, replacement: &Replacement) -> bool {
+    let Replacement::AutoImport { path, name } = replacement else {
+        return false;
+    };
+    let mut import_searcher = ImportSearcher::new(path, name);
+    import_searcher.visit_body(&try_node.body);
+    import_searcher.found_import
+}

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

@@ -1,4 +1,5 @@
 //! Airflow-specific rules.
+pub(crate) mod helpers;
 pub(crate) mod rules;
 
 #[cfg(test)]
@@ -16,6 +17,7 @@ mod tests {
     #[test_case(Rule::AirflowDagNoScheduleArgument, Path::new("AIR002.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_args.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))]
+    #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names_try.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_class_attribute.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_airflow_plugin.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_context.py"))]

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

@@ -1,4 +1,6 @@
 use crate::checkers::ast::Checker;
+use crate::importer::ImportRequest;
+use crate::rules::airflow::helpers::{is_guarded_by_try_except, Replacement};
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::helpers::map_callable;
@@ -59,26 +61,22 @@ impl Violation for Airflow3Removal {
             Replacement::Message(message) => {
                 format!("`{deprecated}` is removed in Airflow 3.0; {message}")
             }
+            Replacement::AutoImport { path: _, name: _ } => {
+                format!("`{deprecated}` is removed in Airflow 3.0")
+            }
         }
     }
 
     fn fix_title(&self) -> Option<String> {
         let Airflow3Removal { replacement, .. } = self;
-        if let Replacement::Name(name) = replacement {
-            Some(format!("Use `{name}` instead"))
-        } else {
-            None
+        match replacement {
+            Replacement::Name(name) => Some(format!("Use `{name}` instead")),
+            Replacement::AutoImport { path, name } => Some(format!("Use `{path}.{name}` instead")),
+            _ => None,
         }
     }
 }
 
-#[derive(Debug, Eq, PartialEq)]
-enum Replacement {
-    None,
-    Name(&'static str),
-    Message(&'static str),
-}
-
 /// AIR302
 pub(crate) fn airflow_3_removal_expr(checker: &Checker, expr: &Expr) {
     if !checker.semantic().seen_module(Modules::AIRFLOW) {
@@ -559,7 +557,9 @@ fn check_method(checker: &Checker, call_expr: &ExprCall) {
 /// SubDagOperator()
 /// ```
 fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
-    let Some(qualified_name) = checker.semantic().resolve_qualified_name(expr) else {
+    let semantic = checker.semantic();
+
+    let Some(qualified_name) = semantic.resolve_qualified_name(expr) else {
         return;
     };
 
@@ -574,7 +574,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
             Replacement::Name("airflow.api_connexion.security.requires_access_*")
         }
         ["airflow", "api_connexion", "security", "requires_access_dataset"] => {
-            Replacement::Name("airflow.api_connexion.security.requires_access_asset")
+            Replacement::AutoImport {
+                path: "airflow.api_connexion.security",
+                name: "requires_access_asset",
+            }
         }
 
         // airflow.auth.managers
@@ -608,9 +611,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         }
 
         // airflow.datasets
-        ["airflow", "Dataset"] | ["airflow", "datasets", "Dataset"] => {
-            Replacement::Name("airflow.sdk.Asset")
-        }
+        ["airflow", "Dataset"] | ["airflow", "datasets", "Dataset"] => Replacement::AutoImport {
+            path: "airflow.sdk",
+            name: "Asset",
+        },
         ["airflow", "datasets", rest @ ..] => match &rest {
             ["DatasetAliasEvent"] => Replacement::None,
             ["DatasetAlias"] => Replacement::Name("airflow.sdk.AssetAlias"),
@@ -904,13 +908,31 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         _ => return,
     };
 
-    checker.report_diagnostic(Diagnostic::new(
+    if is_guarded_by_try_except(expr, &replacement, semantic) {
+        return;
+    }
+
+    let mut diagnostic = Diagnostic::new(
         Airflow3Removal {
             deprecated: qualified_name.to_string(),
-            replacement,
+            replacement: replacement.clone(),
         },
         range,
-    ));
+    );
+
+    if let Replacement::AutoImport { path, name } = replacement {
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer().get_or_import_symbol(
+                &ImportRequest::import_from(path, name),
+                expr.start(),
+                checker.semantic(),
+            )?;
+            let replacement_edit = Edit::range_replacement(binding, range);
+            Ok(Fix::safe_edits(import_edit, [replacement_edit]))
+        });
+    };
+
+    checker.report_diagnostic(diagnostic);
 }
 
 /// Check whether a customized Airflow plugin contains removed extensions.

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

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/airflow/mod.rs
 ---
-AIR302_class_attribute.py:24:21: AIR302 `airflow.Dataset` is removed in Airflow 3.0
+AIR302_class_attribute.py:24:21: AIR302 [*] `airflow.Dataset` is removed in Airflow 3.0
    |
 23 | # airflow.Dataset
 24 | dataset_from_root = DatasetFromRoot()
@@ -11,6 +11,19 @@ AIR302_class_attribute.py:24:21: AIR302 `airflow.Dataset` is removed in Airflow
    |
    = help: Use `airflow.sdk.Asset` instead
 
+ℹ Safe fix
+19 19 | from airflow.providers_manager import ProvidersManager
+20 20 | from airflow.secrets.base_secrets import BaseSecretsBackend
+21 21 | from airflow.secrets.local_filesystem import LocalFilesystemBackend
+   22 |+from airflow.sdk import Asset
+22 23 | 
+23 24 | # airflow.Dataset
+24    |-dataset_from_root = DatasetFromRoot()
+   25 |+dataset_from_root = Asset()
+25 26 | dataset_from_root.iter_datasets()
+26 27 | dataset_from_root.iter_dataset_aliases()
+27 28 | 
+
 AIR302_class_attribute.py:25:19: AIR302 [*] `iter_datasets` is removed in Airflow 3.0
    |
 23 | # airflow.Dataset
@@ -52,7 +65,7 @@ AIR302_class_attribute.py:26:19: AIR302 [*] `iter_dataset_aliases` is removed in
 28 28 | # airflow.datasets
 29 29 | dataset_to_test_method_call = Dataset()
 
-AIR302_class_attribute.py:29:31: AIR302 `airflow.datasets.Dataset` is removed in Airflow 3.0
+AIR302_class_attribute.py:29:31: AIR302 [*] `airflow.datasets.Dataset` is removed in Airflow 3.0
    |
 28 | # airflow.datasets
 29 | dataset_to_test_method_call = Dataset()
@@ -62,6 +75,24 @@ AIR302_class_attribute.py:29:31: AIR302 `airflow.datasets.Dataset` is removed in
    |
    = help: Use `airflow.sdk.Asset` instead
 
+ℹ Safe fix
+19 19 | from airflow.providers_manager import ProvidersManager
+20 20 | from airflow.secrets.base_secrets import BaseSecretsBackend
+21 21 | from airflow.secrets.local_filesystem import LocalFilesystemBackend
+   22 |+from airflow.sdk import Asset
+22 23 | 
+23 24 | # airflow.Dataset
+24 25 | dataset_from_root = DatasetFromRoot()
+--------------------------------------------------------------------------------
+26 27 | dataset_from_root.iter_dataset_aliases()
+27 28 | 
+28 29 | # airflow.datasets
+29    |-dataset_to_test_method_call = Dataset()
+   30 |+dataset_to_test_method_call = Asset()
+30 31 | dataset_to_test_method_call.iter_datasets()
+31 32 | dataset_to_test_method_call.iter_dataset_aliases()
+32 33 | 
+
 AIR302_class_attribute.py:30:29: AIR302 [*] `iter_datasets` is removed in Airflow 3.0
    |
 28 | # airflow.datasets

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

@@ -64,7 +64,7 @@ AIR302_names.py:121:39: AIR302 `airflow.PY312` is removed in Airflow 3.0
     |
     = help: Use `sys.version_info` instead
 
-AIR302_names.py:122:1: AIR302 `airflow.Dataset` is removed in Airflow 3.0
+AIR302_names.py:122:1: AIR302 [*] `airflow.Dataset` is removed in Airflow 3.0
     |
 120 | # airflow root
 121 | PY36, PY37, PY38, PY39, PY310, PY311, PY312
@@ -75,6 +75,20 @@ AIR302_names.py:122:1: AIR302 `airflow.Dataset` is removed in Airflow 3.0
     |
     = help: Use `airflow.sdk.Asset` instead
 
+ℹ Safe fix
+116 116 | from airflow.utils.trigger_rule import TriggerRule
+117 117 | from airflow.www.auth import has_access, has_access_dataset
+118 118 | from airflow.www.utils import get_sensitive_variables_fields, should_hide_value_for_key
+    119 |+from airflow.sdk import Asset
+119 120 | 
+120 121 | # airflow root
+121 122 | PY36, PY37, PY38, PY39, PY310, PY311, PY312
+122     |-DatasetFromRoot()
+    123 |+Asset()
+123 124 | 
+124 125 | # airflow.api_connexion.security
+125 126 | requires_access, requires_access_dataset
+
 AIR302_names.py:125:1: AIR302 `airflow.api_connexion.security.requires_access` is removed in Airflow 3.0
     |
 124 | # airflow.api_connexion.security
@@ -85,7 +99,7 @@ AIR302_names.py:125:1: AIR302 `airflow.api_connexion.security.requires_access` i
     |
     = help: Use `airflow.api_connexion.security.requires_access_*` instead
 
-AIR302_names.py:125:18: AIR302 `airflow.api_connexion.security.requires_access_dataset` is removed in Airflow 3.0
+AIR302_names.py:125:18: AIR302 [*] `airflow.api_connexion.security.requires_access_dataset` is removed in Airflow 3.0
     |
 124 | # airflow.api_connexion.security
 125 | requires_access, requires_access_dataset
@@ -95,6 +109,25 @@ AIR302_names.py:125:18: AIR302 `airflow.api_connexion.security.requires_access_d
     |
     = help: Use `airflow.api_connexion.security.requires_access_asset` instead
 
+ℹ Safe fix
+12  12  | from airflow import (
+13  13  |     Dataset as DatasetFromRoot,
+14  14  | )
+15      |-from airflow.api_connexion.security import requires_access, requires_access_dataset
+    15  |+from airflow.api_connexion.security import requires_access, requires_access_dataset, requires_access_asset
+16  16  | from airflow.auth.managers.base_auth_manager import is_authorized_dataset
+17  17  | from airflow.auth.managers.models.resource_details import DatasetDetails
+18  18  | from airflow.configuration import (
+--------------------------------------------------------------------------------
+122 122 | DatasetFromRoot()
+123 123 | 
+124 124 | # airflow.api_connexion.security
+125     |-requires_access, requires_access_dataset
+    125 |+requires_access, requires_access_asset
+126 126 | 
+127 127 | # airflow.auth.managers
+128 128 | is_authorized_dataset
+
 AIR302_names.py:128:1: AIR302 `airflow.auth.managers.base_auth_manager.is_authorized_dataset` is removed in Airflow 3.0
     |
 127 | # airflow.auth.managers
@@ -188,7 +221,7 @@ AIR302_names.py:136:1: AIR302 `airflow.contrib.aws_athena_hook.AWSAthenaHook` is
 138 | # airflow.datasets
     |
 
-AIR302_names.py:139:1: AIR302 `airflow.datasets.Dataset` is removed in Airflow 3.0
+AIR302_names.py:139:1: AIR302 [*] `airflow.datasets.Dataset` is removed in Airflow 3.0
     |
 138 | # airflow.datasets
 139 | Dataset()
@@ -198,6 +231,24 @@ AIR302_names.py:139:1: AIR302 `airflow.datasets.Dataset` is removed in Airflow 3
     |
     = help: Use `airflow.sdk.Asset` instead
 
+ℹ Safe fix
+116 116 | from airflow.utils.trigger_rule import TriggerRule
+117 117 | from airflow.www.auth import has_access, has_access_dataset
+118 118 | from airflow.www.utils import get_sensitive_variables_fields, should_hide_value_for_key
+    119 |+from airflow.sdk import Asset
+119 120 | 
+120 121 | # airflow root
+121 122 | PY36, PY37, PY38, PY39, PY310, PY311, PY312
+--------------------------------------------------------------------------------
+136 137 | AWSAthenaHook()
+137 138 | 
+138 139 | # airflow.datasets
+139     |-Dataset()
+    140 |+Asset()
+140 141 | DatasetAlias()
+141 142 | DatasetAliasEvent()
+142 143 | DatasetAll()
+
 AIR302_names.py:140:1: AIR302 `airflow.datasets.DatasetAlias` is removed in Airflow 3.0
     |
 138 | # airflow.datasets

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

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

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

@@ -0,0 +1,45 @@
+use ruff_python_ast::statement_visitor::StatementVisitor;
+use ruff_python_ast::{statement_visitor, Alias, Stmt, StmtImportFrom};
+
+/// AST visitor that searches an AST tree for [`ast::StmtImportFrom`] nodes
+/// that match a certain [`QualifiedName`].
+pub(crate) struct ImportSearcher<'a> {
+    module: &'a str,
+    name: &'a str,
+    pub found_import: bool,
+}
+
+impl<'a> ImportSearcher<'a> {
+    pub(crate) fn new(module: &'a str, name: &'a str) -> Self {
+        Self {
+            module,
+            name,
+            found_import: false,
+        }
+    }
+}
+impl StatementVisitor<'_> for ImportSearcher<'_> {
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        if self.found_import {
+            return;
+        }
+        if let Stmt::ImportFrom(StmtImportFrom { module, names, .. }) = stmt {
+            if module.as_ref().is_some_and(|module| module == self.module)
+                && names.iter().any(|Alias { name, .. }| name == self.name)
+            {
+                self.found_import = true;
+                return;
+            }
+        }
+        statement_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_import {
+                return;
+            }
+        }
+    }
+}

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

@@ -1,4 +1,5 @@
 //! NumPy-specific rules.
+pub(crate) mod helpers;
 pub(crate) mod rules;
 
 #[cfg(test)]

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

@@ -1,3 +1,4 @@
+use crate::rules::numpy::helpers::ImportSearcher;
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::name::{QualifiedName, QualifiedNameBuilder};
@@ -892,49 +893,3 @@ fn try_block_contains_undeprecated_import(
     import_searcher.visit_body(&try_node.body);
     import_searcher.found_import
 }
-
-/// AST visitor that searches an AST tree for [`ast::StmtImportFrom`] nodes
-/// that match a certain [`QualifiedName`].
-struct ImportSearcher<'a> {
-    module: &'a str,
-    name: &'a str,
-    found_import: bool,
-}
-
-impl<'a> ImportSearcher<'a> {
-    fn new(module: &'a str, name: &'a str) -> Self {
-        Self {
-            module,
-            name,
-            found_import: false,
-        }
-    }
-}
-
-impl StatementVisitor<'_> for ImportSearcher<'_> {
-    fn visit_stmt(&mut self, stmt: &ast::Stmt) {
-        if self.found_import {
-            return;
-        }
-        if let ast::Stmt::ImportFrom(ast::StmtImportFrom { module, names, .. }) = stmt {
-            if module.as_ref().is_some_and(|module| module == self.module)
-                && names
-                    .iter()
-                    .any(|ast::Alias { name, .. }| name == self.name)
-            {
-                self.found_import = true;
-                return;
-            }
-        }
-        ast::statement_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_import {
-                return;
-            }
-        }
-    }
-}