[airflow] Extend AIR311 and AIR312 rules

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

@@ -74,3 +74,36 @@
 
 # airflow.utils.dag_parsing_context
 get_parsing_context()
+
+from airflow.decorators.base import (
+    DecoratedMappedOperator,
+    DecoratedOperator,
+    TaskDecorator,
+    get_unique_task_id,
+    task_decorator_factory,
+)
+
+# airflow.decorators.base
+DecoratedMappedOperator()
+DecoratedOperator()
+TaskDecorator()
+get_unique_task_id()
+task_decorator_factory()
+
+
+from airflow.models import Param
+
+# airflow.models
+Param()
+
+
+from airflow.sensors.base import (
+    BaseSensorOperator,
+    PokeReturnValue,
+    poke_mode_only,
+)
+
+# airflow.sensors.base
+BaseSensorOperator()
+PokeReturnValue()
+poke_mode_only()

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

@@ -9,7 +9,6 @@
 from airflow.operators.latest_only import LatestOnlyOperator
 from airflow.operators.trigger_dagrun import TriggerDagRunOperator
 from airflow.operators.weekday import BranchDayOfWeekOperator
-from airflow.sensors.date_time import DateTimeSensor
 
 FSHook()
 PackageIndexHook()
@@ -22,24 +21,30 @@
 
 LatestOnlyOperator()
 BranchDayOfWeekOperator()
-DateTimeSensor()
 
 from airflow.operators.python import (
     BranchPythonOperator,
     PythonOperator,
     PythonVirtualenvOperator,
     ShortCircuitOperator,
 )
+from airflow.sensors.bash import BashSensor
+from airflow.sensors.date_time import DateTimeSensor
+
+BranchPythonOperator()
+PythonOperator()
+PythonVirtualenvOperator()
+ShortCircuitOperator()
+
+BashSensor()
+DateTimeSensor()
 from airflow.sensors.date_time import DateTimeSensorAsync
 from airflow.sensors.external_task import (
     ExternalTaskMarker,
     ExternalTaskSensor,
 )
-from airflow.sensors.time_sensor import (
-    TimeSensor,
-    TimeSensorAsync,
-)
 from airflow.sensors.filesystem import FileSensor
+from airflow.sensors.python import PythonSensor
 
 BranchPythonOperator()
 PythonOperator()
@@ -49,6 +54,13 @@
 ExternalTaskMarker()
 ExternalTaskSensor()
 FileSensor()
+PythonSensor()
+
+from airflow.sensors.time_sensor import (
+    TimeSensor,
+    TimeSensorAsync,
+)
+
 TimeSensor()
 TimeSensorAsync()
 

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

@@ -215,6 +215,12 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
                 version: "0.0.1",
             }
         }
+        ["airflow", "sensors", "bash", "BashSensor"] => ProviderReplacement::AutoImport {
+            module: "airflow.providers.standard.sensor.bash",
+            name: "BashSensor",
+            provider: "standard",
+            version: "0.0.1",
+        },
         [
             "airflow",
             "sensors",
@@ -243,6 +249,12 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
             provider: "standard",
             version: "0.0.2",
         },
+        ["airflow", "sensors", "python", "PythonSensor"] => ProviderReplacement::AutoImport {
+            module: "airflow.providers.standard.sensors.python",
+            name: "PythonSensor",
+            provider: "standard",
+            version: "0.0.1",
+        },
         [
             "airflow",
             "sensors",

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

@@ -227,6 +227,19 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
             module: "airflow.sdk",
             name: (*rest).to_string(),
         },
+        [
+            "airflow",
+            "decorators",
+            "base",
+            rest @ ("DecoratedMappedOperator"
+            | "DecoratedOperator"
+            | "TaskDecorator"
+            | "get_unique_task_id"
+            | "task_decorator_factory"),
+        ] => Replacement::SourceModuleMoved {
+            module: "airflow.sdk.bases.decorator",
+            name: (*rest).to_string(),
+        },
 
         // airflow.io
         ["airflow", "io", "path", "ObjectStoragePath"] => Replacement::SourceModuleMoved {
@@ -245,6 +258,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
                 name: (*rest).to_string(),
             }
         }
+        ["airflow", "models", "Param"] => Replacement::AutoImport {
+            module: "airflow.sdk.definitions.param",
+            name: "Param",
+        },
 
         // airflow.models.baseoperator
         [
@@ -260,16 +277,30 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
             module: "airflow.sdk",
             name: "BaseOperatorLink",
         },
+
         // airflow.model..DAG
         ["airflow", "models", .., "DAG"] => Replacement::SourceModuleMoved {
             module: "airflow.sdk",
             name: "DAG".to_string(),
         },
+
+        // airflow.sensors.base
+        [
+            "airflow",
+            "sensors",
+            "base",
+            rest @ ("BaseSensorOperator" | "PokeReturnValue" | "poke_mode_only"),
+        ] => Replacement::SourceModuleMoved {
+            module: "airflow.sdk",
+            name: (*rest).to_string(),
+        },
+
         // airflow.timetables
         ["airflow", "timetables", "datasets", "DatasetOrTimeSchedule"] => Replacement::AutoImport {
             module: "airflow.timetables.assets",
             name: "AssetOrTimeSchedule",
         },
+
         // airflow.utils
         [
             "airflow",