Skip to content

Commit 5663426

Browse files
authored
[airflow] Extend AIR311 and AIR312 rules (#20082)
<!-- 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? --> Extend the following rules. ### AIR311 * `airflow.sensors.base.BaseSensorOperator` → airflow.sdk.bases.sensor.BaseSensorOperator` * `airflow.sensors.base.PokeReturnValue` → airflow.sdk.bases.sensor.PokeReturnValue` * `airflow.sensors.base.poke_mode_only` → airflow.sdk.bases.sensor.poke_mode_only` * `airflow.decorators.base.DecoratedOperator` → airflow.sdk.bases.decorator.DecoratedOperator` * `airflow.models.param.Param` → airflow.sdk.definitions.param.Param` * `airflow.decorators.base.DecoratedMappedOperator` → `airflow.sdk.bases.decorator.DecoratedMappedOperator` * `airflow.decorators.base.DecoratedOperator` → `airflow.sdk.bases.decorator.DecoratedOperator` * `airflow.decorators.base.TaskDecorator` → `airflow.sdk.bases.decorator.TaskDecorator` * `airflow.decorators.base.get_unique_task_id` → `airflow.sdk.bases.decorator.get_unique_task_id` * `airflow.decorators.base.task_decorator_factory` → `airflow.sdk.bases.decorator.task_decorator_factory` ### AIR312 * `airflow.sensors.bash.BashSensor` → `airflow.providers.standard.sensor.bash.BashSensor` * `airflow.sensors.python.PythonSensor` → `airflow.providers.standard.sensors.python.PythonSensor` ## Test Plan <!-- How was it tested? --> update the test fixture accordingly in the second commit and reorg in the third
1 parent 0b35487 commit 5663426

File tree

6 files changed

+875
-446
lines changed

6 files changed

+875
-446
lines changed

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,3 +74,36 @@
7474

7575
# airflow.utils.dag_parsing_context
7676
get_parsing_context()
77+
78+
from airflow.decorators.base import (
79+
DecoratedMappedOperator,
80+
DecoratedOperator,
81+
TaskDecorator,
82+
get_unique_task_id,
83+
task_decorator_factory,
84+
)
85+
86+
# airflow.decorators.base
87+
DecoratedMappedOperator()
88+
DecoratedOperator()
89+
TaskDecorator()
90+
get_unique_task_id()
91+
task_decorator_factory()
92+
93+
94+
from airflow.models import Param
95+
96+
# airflow.models
97+
Param()
98+
99+
100+
from airflow.sensors.base import (
101+
BaseSensorOperator,
102+
PokeReturnValue,
103+
poke_mode_only,
104+
)
105+
106+
# airflow.sensors.base
107+
BaseSensorOperator()
108+
PokeReturnValue()
109+
poke_mode_only()

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

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
from airflow.operators.latest_only import LatestOnlyOperator
1010
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
1111
from airflow.operators.weekday import BranchDayOfWeekOperator
12-
from airflow.sensors.date_time import DateTimeSensor
1312

1413
FSHook()
1514
PackageIndexHook()
@@ -22,24 +21,30 @@
2221

2322
LatestOnlyOperator()
2423
BranchDayOfWeekOperator()
25-
DateTimeSensor()
2624

2725
from airflow.operators.python import (
2826
BranchPythonOperator,
2927
PythonOperator,
3028
PythonVirtualenvOperator,
3129
ShortCircuitOperator,
3230
)
31+
from airflow.sensors.bash import BashSensor
32+
from airflow.sensors.date_time import DateTimeSensor
33+
34+
BranchPythonOperator()
35+
PythonOperator()
36+
PythonVirtualenvOperator()
37+
ShortCircuitOperator()
38+
39+
BashSensor()
40+
DateTimeSensor()
3341
from airflow.sensors.date_time import DateTimeSensorAsync
3442
from airflow.sensors.external_task import (
3543
ExternalTaskMarker,
3644
ExternalTaskSensor,
3745
)
38-
from airflow.sensors.time_sensor import (
39-
TimeSensor,
40-
TimeSensorAsync,
41-
)
4246
from airflow.sensors.filesystem import FileSensor
47+
from airflow.sensors.python import PythonSensor
4348

4449
BranchPythonOperator()
4550
PythonOperator()
@@ -49,6 +54,13 @@
4954
ExternalTaskMarker()
5055
ExternalTaskSensor()
5156
FileSensor()
57+
PythonSensor()
58+
59+
from airflow.sensors.time_sensor import (
60+
TimeSensor,
61+
TimeSensorAsync,
62+
)
63+
5264
TimeSensor()
5365
TimeSensorAsync()
5466

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,12 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
215215
version: "0.0.1",
216216
}
217217
}
218+
["airflow", "sensors", "bash", "BashSensor"] => ProviderReplacement::AutoImport {
219+
module: "airflow.providers.standard.sensor.bash",
220+
name: "BashSensor",
221+
provider: "standard",
222+
version: "0.0.1",
223+
},
218224
[
219225
"airflow",
220226
"sensors",
@@ -243,6 +249,12 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
243249
provider: "standard",
244250
version: "0.0.2",
245251
},
252+
["airflow", "sensors", "python", "PythonSensor"] => ProviderReplacement::AutoImport {
253+
module: "airflow.providers.standard.sensors.python",
254+
name: "PythonSensor",
255+
provider: "standard",
256+
version: "0.0.1",
257+
},
246258
[
247259
"airflow",
248260
"sensors",

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,19 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
227227
module: "airflow.sdk",
228228
name: (*rest).to_string(),
229229
},
230+
[
231+
"airflow",
232+
"decorators",
233+
"base",
234+
rest @ ("DecoratedMappedOperator"
235+
| "DecoratedOperator"
236+
| "TaskDecorator"
237+
| "get_unique_task_id"
238+
| "task_decorator_factory"),
239+
] => Replacement::SourceModuleMoved {
240+
module: "airflow.sdk.bases.decorator",
241+
name: (*rest).to_string(),
242+
},
230243

231244
// airflow.io
232245
["airflow", "io", "path", "ObjectStoragePath"] => Replacement::SourceModuleMoved {
@@ -245,6 +258,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
245258
name: (*rest).to_string(),
246259
}
247260
}
261+
["airflow", "models", "Param"] => Replacement::AutoImport {
262+
module: "airflow.sdk.definitions.param",
263+
name: "Param",
264+
},
248265

249266
// airflow.models.baseoperator
250267
[
@@ -260,16 +277,30 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
260277
module: "airflow.sdk",
261278
name: "BaseOperatorLink",
262279
},
280+
263281
// airflow.model..DAG
264282
["airflow", "models", .., "DAG"] => Replacement::SourceModuleMoved {
265283
module: "airflow.sdk",
266284
name: "DAG".to_string(),
267285
},
286+
287+
// airflow.sensors.base
288+
[
289+
"airflow",
290+
"sensors",
291+
"base",
292+
rest @ ("BaseSensorOperator" | "PokeReturnValue" | "poke_mode_only"),
293+
] => Replacement::SourceModuleMoved {
294+
module: "airflow.sdk",
295+
name: (*rest).to_string(),
296+
},
297+
268298
// airflow.timetables
269299
["airflow", "timetables", "datasets", "DatasetOrTimeSchedule"] => Replacement::AutoImport {
270300
module: "airflow.timetables.assets",
271301
name: "AssetOrTimeSchedule",
272302
},
303+
273304
// airflow.utils
274305
[
275306
"airflow",

0 commit comments

Comments
 (0)