feat(AIR311): apply Replacement::AutoImport

Author: Lee-W

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

@@ -194,20 +194,33 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
 
     let replacement = match qualified_name.segments() {
         // airflow.datasets.metadata
-        ["airflow", "datasets", "metadata", "Metadata"] => {
-            Replacement::Name("airflow.sdk.Metadata")
-        }
+        ["airflow", "datasets", "metadata", "Metadata"] => Replacement::AutoImport {
+            module: "airflow.sdk",
+            name: "Metadata",
+        },
         // airflow.datasets
         ["airflow", "Dataset"] | ["airflow", "datasets", "Dataset"] => Replacement::AutoImport {
             module: "airflow.sdk",
             name: "Asset",
         },
         ["airflow", "datasets", rest] => match *rest {
             "DatasetAliasEvent" => Replacement::None,
-            "DatasetAlias" => Replacement::Name("airflow.sdk.AssetAlias"),
-            "DatasetAll" => Replacement::Name("airflow.sdk.AssetAll"),
-            "DatasetAny" => Replacement::Name("airflow.sdk.AssetAny"),
-            "expand_alias_to_datasets" => Replacement::Name("airflow.sdk.expand_alias_to_assets"),
+            "DatasetAlias" => Replacement::AutoImport {
+                module: "airflow.sdk",
+                name: "AssetAlias",
+            },
+            "DatasetAll" => Replacement::AutoImport {
+                module: "airflow.sdk",
+                name: "AssetAll",
+            },
+            "DatasetAny" => Replacement::AutoImport {
+                module: "airflow.sdk",
+                name: "AssetAny",
+            },
+            "expand_alias_to_datasets" => Replacement::AutoImport {
+                module: "airflow.sdk",
+                name: "expand_alias_to_assets",
+            },
             _ => return,
         },
 
@@ -235,9 +248,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
                 module: "airflow.sdk",
                 name: (*rest).to_string(),
             },
-            "BaseOperatorLink" => {
-                Replacement::Name("airflow.sdk.definitions.baseoperatorlink.BaseOperatorLink")
-            }
+            "BaseOperatorLink" => Replacement::AutoImport {
+                module: "airflow.sdk.definitions.baseoperatorlink",
+                name: "BaseOperatorLink",
+            },
             _ => return,
         },
         // airflow.model..DAG
@@ -246,12 +260,16 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
             name: "DAG".to_string(),
         },
         // airflow.timetables
-        ["airflow", "timetables", "datasets", "DatasetOrTimeSchedule"] => {
-            Replacement::Name("airflow.timetables.assets.AssetOrTimeSchedule")
-        }
+        ["airflow", "timetables", "datasets", "DatasetOrTimeSchedule"] => Replacement::AutoImport {
+            module: "airflow.timetables.assets",
+            name: "AssetOrTimeSchedule",
+        },
         // airflow.utils
         ["airflow", "utils", "dag_parsing_context", "get_parsing_context"] => {
-            Replacement::Name("airflow.sdk.get_parsing_context")
+            Replacement::AutoImport {
+                module: "airflow.sdk",
+                name: "get_parsing_context",
+            }
         }
 
         _ => return,

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

@@ -50,7 +50,7 @@ AIR311_names.py:26:1: AIR311 [*] `airflow.datasets.Dataset` is removed in Airflo
 28 29 | DatasetAll()
 29 30 | DatasetAny()
 
-AIR311_names.py:27:1: AIR311 `airflow.datasets.DatasetAlias` 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:27:1: AIR311 [*] `airflow.datasets.DatasetAlias` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 25 | # airflow.datasets
 26 | Dataset()
@@ -61,7 +61,24 @@ AIR311_names.py:27:1: AIR311 `airflow.datasets.DatasetAlias` is removed in Airfl
    |
    = help: Use `airflow.sdk.AssetAlias` instead
 
-AIR311_names.py:28:1: AIR311 `airflow.datasets.DatasetAll` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+ℹ Safe fix
+18 18 | from airflow.models.dag import DAG as DAGFromDag
+19 19 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+20 20 | from airflow.utils.dag_parsing_context import get_parsing_context
+   21 |+from airflow.sdk import AssetAlias
+21 22 | 
+22 23 | # airflow
+23 24 | DatasetFromRoot()
+24 25 | 
+25 26 | # airflow.datasets
+26 27 | Dataset()
+27    |-DatasetAlias()
+   28 |+AssetAlias()
+28 29 | DatasetAll()
+29 30 | DatasetAny()
+30 31 | Metadata()
+
+AIR311_names.py:28:1: AIR311 [*] `airflow.datasets.DatasetAll` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 26 | Dataset()
 27 | DatasetAlias()
@@ -72,7 +89,25 @@ AIR311_names.py:28:1: AIR311 `airflow.datasets.DatasetAll` is removed in Airflow
    |
    = help: Use `airflow.sdk.AssetAll` instead
 
-AIR311_names.py:29:1: AIR311 `airflow.datasets.DatasetAny` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+ℹ Safe fix
+18 18 | from airflow.models.dag import DAG as DAGFromDag
+19 19 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+20 20 | from airflow.utils.dag_parsing_context import get_parsing_context
+   21 |+from airflow.sdk import AssetAll
+21 22 | 
+22 23 | # airflow
+23 24 | DatasetFromRoot()
+--------------------------------------------------------------------------------
+25 26 | # airflow.datasets
+26 27 | Dataset()
+27 28 | DatasetAlias()
+28    |-DatasetAll()
+   29 |+AssetAll()
+29 30 | DatasetAny()
+30 31 | Metadata()
+31 32 | expand_alias_to_datasets()
+
+AIR311_names.py:29:1: AIR311 [*] `airflow.datasets.DatasetAny` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 27 | DatasetAlias()
 28 | DatasetAll()
@@ -83,6 +118,24 @@ AIR311_names.py:29:1: AIR311 `airflow.datasets.DatasetAny` is removed in Airflow
    |
    = help: Use `airflow.sdk.AssetAny` instead
 
+ℹ Safe fix
+18 18 | from airflow.models.dag import DAG as DAGFromDag
+19 19 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+20 20 | from airflow.utils.dag_parsing_context import get_parsing_context
+   21 |+from airflow.sdk import AssetAny
+21 22 | 
+22 23 | # airflow
+23 24 | DatasetFromRoot()
+--------------------------------------------------------------------------------
+26 27 | Dataset()
+27 28 | DatasetAlias()
+28 29 | DatasetAll()
+29    |-DatasetAny()
+   30 |+AssetAny()
+30 31 | Metadata()
+31 32 | expand_alias_to_datasets()
+32 33 | 
+
 AIR311_names.py:30:1: AIR311 `airflow.datasets.metadata.Metadata` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 28 | DatasetAll()
@@ -93,7 +146,7 @@ AIR311_names.py:30:1: AIR311 `airflow.datasets.metadata.Metadata` is removed in
    |
    = help: Use `airflow.sdk.Metadata` instead
 
-AIR311_names.py:31:1: AIR311 `airflow.datasets.expand_alias_to_datasets` 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:31:1: AIR311 [*] `airflow.datasets.expand_alias_to_datasets` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 29 | DatasetAny()
 30 | Metadata()
@@ -104,6 +157,24 @@ AIR311_names.py:31:1: AIR311 `airflow.datasets.expand_alias_to_datasets` is remo
    |
    = help: Use `airflow.sdk.expand_alias_to_assets` instead
 
+ℹ Safe fix
+18 18 | from airflow.models.dag import DAG as DAGFromDag
+19 19 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+20 20 | from airflow.utils.dag_parsing_context import get_parsing_context
+   21 |+from airflow.sdk import expand_alias_to_assets
+21 22 | 
+22 23 | # airflow
+23 24 | DatasetFromRoot()
+--------------------------------------------------------------------------------
+28 29 | DatasetAll()
+29 30 | DatasetAny()
+30 31 | Metadata()
+31    |-expand_alias_to_datasets()
+   32 |+expand_alias_to_assets()
+32 33 | 
+33 34 | # airflow.decorators
+34 35 | dag()
+
 AIR311_names.py:34:1: AIR311 `airflow.decorators.dag` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 33 | # airflow.decorators
@@ -228,7 +299,7 @@ AIR311_names.py:56:1: AIR311 `airflow.models.dag.DAG` is removed in Airflow 3.0;
    |
    = help: Use `airflow.sdk.DAG` instead
 
-AIR311_names.py:58: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.
+AIR311_names.py:58: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.
    |
 56 | DAGFromDag()
 57 | # airflow.timetables.datasets
@@ -239,6 +310,24 @@ AIR311_names.py:58:1: AIR311 `airflow.timetables.datasets.DatasetOrTimeSchedule`
    |
    = help: Use `airflow.timetables.assets.AssetOrTimeSchedule` instead
 
+ℹ Safe fix
+18 18 | from airflow.models.dag import DAG as DAGFromDag
+19 19 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+20 20 | from airflow.utils.dag_parsing_context import get_parsing_context
+   21 |+from airflow.timetables.assets import AssetOrTimeSchedule
+21 22 | 
+22 23 | # airflow
+23 24 | DatasetFromRoot()
+--------------------------------------------------------------------------------
+55 56 | # airflow.models.dag
+56 57 | DAGFromDag()
+57 58 | # airflow.timetables.datasets
+58    |-DatasetOrTimeSchedule()
+   59 |+AssetOrTimeSchedule()
+59 60 | 
+60 61 | # airflow.utils.dag_parsing_context
+61 62 | get_parsing_context()
+
 AIR311_names.py:61:1: AIR311 `airflow.utils.dag_parsing_context.get_parsing_context` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
 60 | # airflow.utils.dag_parsing_context