[airflow] Add unsafe fix for module moved cases (AIR311) (#18366)

<!--
Thank you for contributing to Ruff/ty! 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? (Please prefix
with `[ty]` for ty pull
  requests.)
- Does this pull request include references to any relevant issues?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->

Follow up on #18093 and apply it
to AIR311

---

Rules fixed
* `airflow.models.datasets.expand_alias_to_datasets` →
`airflow.models.asset.expand_alias_to_assets`
* `airflow.models.baseoperatorlink.BaseOperatorLink` →
`airflow.sdk.BaseOperatorLink`


## Test Plan

<!-- How was it tested? -->
The existing test fixtures have been updated

Author: Lee-W

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

@@ -9,19 +9,12 @@
     expand_alias_to_datasets,
 )
 from airflow.datasets.metadata import Metadata
-from airflow.decorators import dag, setup, task, task_group, teardown
-from airflow.io.path import ObjectStoragePath
-from airflow.io.storage import attach
-from airflow.models import DAG as DAGFromModel
-from airflow.models import (
-    Connection,
-    Variable,
+from airflow.decorators import (
+    dag,
+    setup,
+    task,
+    task_group,
 )
-from airflow.models.baseoperator import chain, chain_linear, cross_downstream
-from airflow.models.baseoperatorlink import BaseOperatorLink
-from airflow.models.dag import DAG as DAGFromDag
-from airflow.timetables.datasets import DatasetOrTimeSchedule
-from airflow.utils.dag_parsing_context import get_parsing_context
 
 # airflow
 DatasetFromRoot()
@@ -39,9 +32,22 @@
 task()
 task_group()
 setup()
+from airflow.decorators import teardown
+from airflow.io.path import ObjectStoragePath
+from airflow.io.storage import attach
+from airflow.models import DAG as DAGFromModel
+from airflow.models import (
+    Connection,
+    Variable,
+)
+from airflow.models.baseoperator import chain, chain_linear, cross_downstream
+from airflow.models.baseoperatorlink import BaseOperatorLink
+from airflow.models.dag import DAG as DAGFromDag
+
+# airflow.decorators
 teardown()
 
-# airflow.io
+# # airflow.io
 ObjectStoragePath()
 attach()
 
@@ -60,6 +66,9 @@
 
 # airflow.models.dag
 DAGFromDag()
+from airflow.timetables.datasets import DatasetOrTimeSchedule
+from airflow.utils.dag_parsing_context import get_parsing_context
+
 # airflow.timetables.datasets
 DatasetOrTimeSchedule()
 

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

@@ -1,7 +1,7 @@
 use crate::checkers::ast::Checker;
-use crate::importer::ImportRequest;
+use crate::rules::airflow::helpers::{Replacement, is_airflow_builtin_or_provider};
 use crate::rules::airflow::helpers::{
-    Replacement, is_airflow_builtin_or_provider, is_guarded_by_try_except,
+    generate_import_edit, generate_remove_and_runtime_import_edit, is_guarded_by_try_except,
 };
 use crate::{Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{ViolationMetadata, derive_message_formats};
@@ -211,7 +211,7 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
                 name: "AssetAny",
             },
             "expand_alias_to_datasets" => Replacement::AutoImport {
-                module: "airflow.sdk",
+                module: "airflow.models.asset",
                 name: "expand_alias_to_assets",
             },
             _ => return,
@@ -256,7 +256,7 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
             name: (*rest).to_string(),
         },
         ["airflow", "models", "baseoperatorlink", "BaseOperatorLink"] => Replacement::AutoImport {
-            module: "airflow.sdk.definitions.baseoperatorlink",
+            module: "airflow.sdk",
             name: "BaseOperatorLink",
         },
         // airflow.model..DAG
@@ -301,22 +301,16 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
     if is_guarded_by_try_except(expr, module, name, checker.semantic()) {
         return;
     }
-
-    checker
-        .report_diagnostic(
-            Airflow3SuggestedUpdate {
-                deprecated: qualified_name.to_string(),
-                replacement: replacement.clone(),
-            },
-            range,
-        )
-        .try_set_fix(|| {
-            let (import_edit, binding) = checker.importer().get_or_import_symbol(
-                &ImportRequest::import_from(module, name),
-                expr.start(),
-                checker.semantic(),
-            )?;
-            let replacement_edit = Edit::range_replacement(binding, range);
-            Ok(Fix::safe_edits(import_edit, [replacement_edit]))
-        });
+    let mut diagnostic = checker.report_diagnostic(
+        Airflow3SuggestedUpdate {
+            deprecated: qualified_name.to_string(),
+            replacement: replacement.clone(),
+        },
+        range,
+    );
+    if let Some(fix) = generate_import_edit(expr, checker, module, name, range)
+        .or_else(|| generate_remove_and_runtime_import_edit(expr, checker, module, name))
+    {
+        diagnostic.set_fix(fix);
+    }
 }