[airflow] Add unsafe fix for module moved cases (AIR301) (#18367)

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

## Test Plan

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

The existing test fixtures have been updated

Author: Lee-W

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

@@ -10,22 +10,10 @@
     PY312,
 )
 from airflow.api_connexion.security import requires_access
-from airflow.configuration import (
-    as_dict,
-    get,
-    getboolean,
-    getfloat,
-    getint,
-    has_option,
-    remove_option,
-    set,
-)
 from airflow.contrib.aws_athena_hook import AWSAthenaHook
 from airflow.datasets import DatasetAliasEvent
-from airflow.hooks.base_hook import BaseHook
 from airflow.operators.subdag import SubDagOperator
 from airflow.secrets.local_filesystem import LocalFilesystemBackend
-from airflow.sensors.base_sensor_operator import BaseSensorOperator
 from airflow.triggers.external_task import TaskStateTrigger
 from airflow.utils import dates
 from airflow.utils.dag_cycle_tester import test_cycle
@@ -40,13 +28,10 @@
 )
 from airflow.utils.db import create_session
 from airflow.utils.decorators import apply_defaults
-from airflow.utils.file import TemporaryDirectory, mkdirs
-from airflow.utils.helpers import chain as helper_chain
-from airflow.utils.helpers import cross_downstream as helper_cross_downstream
-from airflow.utils.log import secrets_masker
+from airflow.utils.file import mkdirs
 from airflow.utils.state import SHUTDOWN, terminating_states
 from airflow.utils.trigger_rule import TriggerRule
-from airflow.www.auth import has_access
+from airflow.www.auth import has_access, has_access_dataset
 from airflow.www.utils import get_sensitive_variables_fields, should_hide_value_for_key
 
 # airflow root
@@ -55,11 +40,6 @@
 # airflow.api_connexion.security
 requires_access
 
-
-# airflow.configuration
-get, getboolean, getfloat, getint, has_option, remove_option, as_dict, set
-
-
 # airflow.contrib.*
 AWSAthenaHook()
 
@@ -68,10 +48,6 @@
 DatasetAliasEvent()
 
 
-# airflow.hooks
-BaseHook()
-
-
 # airflow.operators.subdag.*
 SubDagOperator()
 
@@ -81,10 +57,6 @@
 LocalFilesystemBackend()
 
 
-# airflow.sensors.base_sensor_operator
-BaseSensorOperator()
-
-
 # airflow.triggers.external_task
 TaskStateTrigger()
 
@@ -114,15 +86,8 @@
 apply_defaults
 
 # airflow.utils.file
-TemporaryDirectory()
 mkdirs
 
-#  airflow.utils.helpers
-helper_chain
-helper_cross_downstream
-
-#  airflow.utils.log
-secrets_masker
 
 # airflow.utils.state
 SHUTDOWN
@@ -135,37 +100,8 @@
 
 # airflow.www.auth
 has_access
+has_access_dataset
 
 # airflow.www.utils
 get_sensitive_variables_fields
 should_hide_value_for_key
-
-# airflow.operators.python
-from airflow.operators.python import get_current_context
-
-get_current_context()
-
-# airflow.providers.mysql
-from airflow.providers.mysql.datasets.mysql import sanitize_uri
-
-sanitize_uri
-
-# airflow.providers.postgres
-from airflow.providers.postgres.datasets.postgres import sanitize_uri
-
-sanitize_uri
-
-# airflow.providers.trino
-from airflow.providers.trino.datasets.trino import sanitize_uri
-
-sanitize_uri
-
-# airflow.notifications.basenotifier
-from airflow.notifications.basenotifier import BaseNotifier
-
-BaseNotifier()
-
-# airflow.auth.manager
-from airflow.auth.managers.base_auth_manager import BaseAuthManager
-
-BaseAuthManager()

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

@@ -3,7 +3,6 @@
 from airflow.api_connexion.security import requires_access_dataset
 from airflow.auth.managers.models.resource_details import (
     DatasetDetails,
-
 )
 from airflow.datasets.manager import (
     DatasetManager,
@@ -12,15 +11,13 @@
 )
 from airflow.lineage.hook import DatasetLineageInfo
 from airflow.metrics.validators import AllowListValidator, BlockListValidator
-from airflow.secrets.local_filesystm import load_connections
+from airflow.secrets.local_filesystem import load_connections
 from airflow.security.permissions import RESOURCE_DATASET
-from airflow.www.auth import has_access_dataset
 
 requires_access_dataset()
 
 DatasetDetails()
 
-
 DatasetManager()
 dataset_manager()
 resolve_dataset_manager()
@@ -34,7 +31,6 @@
 
 RESOURCE_DATASET
 
-has_access_dataset()
 
 from airflow.listeners.spec.dataset import (
     on_dataset_changed,
@@ -43,3 +39,76 @@
 
 on_dataset_created()
 on_dataset_changed()
+
+
+# airflow.operators.python
+from airflow.operators.python import get_current_context
+
+get_current_context()
+
+# airflow.providers.mysql
+from airflow.providers.mysql.datasets.mysql import sanitize_uri
+
+sanitize_uri
+
+# airflow.providers.postgres
+from airflow.providers.postgres.datasets.postgres import sanitize_uri
+
+sanitize_uri
+
+# airflow.providers.trino
+from airflow.providers.trino.datasets.trino import sanitize_uri
+
+sanitize_uri
+
+# airflow.notifications.basenotifier
+from airflow.notifications.basenotifier import BaseNotifier
+
+BaseNotifier()
+
+# airflow.auth.manager
+from airflow.auth.managers.base_auth_manager import BaseAuthManager
+
+BaseAuthManager()
+
+
+from airflow.configuration import (
+    as_dict,
+    get,
+    getboolean,
+    getfloat,
+    getint,
+    has_option,
+    remove_option,
+    set,
+)
+
+# airflow.configuration
+get, getboolean, getfloat, getint, has_option, remove_option, as_dict, set
+from airflow.hooks.base_hook import BaseHook
+
+# airflow.hooks
+BaseHook()
+
+from airflow.sensors.base_sensor_operator import BaseSensorOperator
+
+# airflow.sensors.base_sensor_operator
+BaseSensorOperator()
+BaseHook()
+
+from airflow.utils.helpers import chain as helper_chain
+from airflow.utils.helpers import cross_downstream as helper_cross_downstream
+
+#  airflow.utils.helpers
+helper_chain
+helper_cross_downstream
+
+# airflow.utils.file
+from airflow.utils.file import TemporaryDirectory
+
+TemporaryDirectory()
+
+from airflow.utils.log import secrets_masker
+
+#  airflow.utils.log
+secrets_masker

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

@@ -1,54 +1,54 @@
 from __future__ import annotations
 
 from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities
-from airflow.providers.amazon.aws.datasets.s3 import (
-    convert_dataset_to_openlineage as s3_convert_dataset_to_openlineage,
-)
-from airflow.providers.amazon.aws.datasets.s3 import create_dataset as s3_create_dataset
-from airflow.providers.common.io.dataset.file import (
-    convert_dataset_to_openlineage as io_convert_dataset_to_openlineage,
-)
-from airflow.providers.common.io.dataset.file import create_dataset as io_create_dataset
-
-from airflow.providers.google.datasets.bigquery import (
-    create_dataset as bigquery_create_dataset,
-)
-from airflow.providers.google.datasets.gcs import (
-    convert_dataset_to_openlineage as gcs_convert_dataset_to_openlineage,
-)
-from airflow.providers.google.datasets.gcs import create_dataset as gcs_create_dataset
 from airflow.providers.openlineage.utils.utils import (
     DatasetInfo,
     translate_airflow_dataset,
 )
+from airflow.secrets.local_filesystem import load_connections
+from airflow.security.permissions import RESOURCE_DATASET
 
 AvpEntities.DATASET
 
+# airflow.providers.openlineage.utils.utils
+DatasetInfo()
+translate_airflow_dataset()
+
+# airflow.secrets.local_filesystem
+load_connections()
+
+# airflow.security.permissions
+RESOURCE_DATASET
+
+from airflow.providers.amazon.aws.datasets.s3 import (
+    convert_dataset_to_openlineage as s3_convert_dataset_to_openlineage,
+)
+from airflow.providers.amazon.aws.datasets.s3 import create_dataset as s3_create_dataset
+
 s3_create_dataset()
 s3_convert_dataset_to_openlineage()
 
+from airflow.providers.common.io.dataset.file import (
+    convert_dataset_to_openlineage as io_convert_dataset_to_openlineage,
+)
+from airflow.providers.common.io.dataset.file import create_dataset as io_create_dataset
+
 io_create_dataset()
 io_convert_dataset_to_openlineage()
 
 
+# # airflow.providers.google.datasets.bigquery
+from airflow.providers.google.datasets.bigquery import (
+    create_dataset as bigquery_create_dataset,
+)
 
-# airflow.providers.google.datasets.bigquery
 bigquery_create_dataset()
+
 # airflow.providers.google.datasets.gcs
+from airflow.providers.google.datasets.gcs import (
+    convert_dataset_to_openlineage as gcs_convert_dataset_to_openlineage,
+)
+from airflow.providers.google.datasets.gcs import create_dataset as gcs_create_dataset
+
 gcs_create_dataset()
 gcs_convert_dataset_to_openlineage()
-# airflow.providers.openlineage.utils.utils
-DatasetInfo()
-translate_airflow_dataset()
-#
-# airflow.secrets.local_filesystem
-load_connections()
-#
-# airflow.security.permissions
-RESOURCE_DATASET
-
-# airflow.timetables
-DatasetTriggeredTimetable()
-#
-# airflow.www.auth
-has_access_dataset

crates/ruff_linter/src/rules/airflow/rules/removal_in_3.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, is_guarded_by_try_except,
+    Replacement, generate_import_edit, generate_remove_and_runtime_import_edit,
+    is_airflow_builtin_or_provider, is_guarded_by_try_except,
 };
 use crate::{Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{ViolationMetadata, derive_message_formats};
@@ -614,7 +614,6 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
         },
 
         // airflow.configuration
-        // TODO: check whether we could improve it
         [
             "airflow",
             "configuration",
@@ -984,24 +983,19 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
     }
 
     let import_target = name.split('.').next().unwrap_or(name);
+    let mut diagnostic = checker.report_diagnostic(
+        Airflow3Removal {
+            deprecated: qualified_name.to_string(),
+            replacement: replacement.clone(),
+        },
+        range,
+    );
 
-    checker
-        .report_diagnostic(
-            Airflow3Removal {
-                deprecated: qualified_name.to_string(),
-                replacement: replacement.clone(),
-            },
-            range,
-        )
-        .try_set_fix(|| {
-            let (import_edit, _) = checker.importer().get_or_import_symbol(
-                &ImportRequest::import_from(module, import_target),
-                expr.start(),
-                checker.semantic(),
-            )?;
-            let replacement_edit = Edit::range_replacement(name.to_string(), range);
-            Ok(Fix::safe_edits(import_edit, [replacement_edit]))
-        });
+    if let Some(fix) = generate_import_edit(expr, checker, module, import_target, range)
+        .or_else(|| generate_remove_and_runtime_import_edit(expr, checker, module, name))
+    {
+        diagnostic.set_fix(fix);
+    }
 }
 
 /// Check whether a customized Airflow plugin contains removed extensions.