Skip to content

Commit 1a79722

Browse files
authored
[airflow] Extend AIR311 rules (#17422)
<!-- Thank you for contributing to Ruff! 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? - 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 AIR311 rules * `airflow.io.path.ObjectStoragePath` → `airflow.sdk.ObjectStoragePath` * `airflow.io.storage.attach` → `airflow.sdk.io.attach` * `airflow.models.dag.DAG` → `airflow.sdk.DAG` * `airflow.models.DAG` → `airflow.sdk.DAG` * `airflow.decorators.dag` → `airflow.sdk.dag` * `airflow.decorators.task` → `airflow.sdk.task` * `airflow.decorators.task_group` → `airflow.sdk.task_group` * `airflow.decorators.setup` → `airflow.sdk.setup` * `airflow.decorators.teardown` → `airflow.sdk.teardown` ## Test Plan <!-- How was it tested? --> The test case has been added to the button of the existing test fixtures, confirmed to be correct and later reorgnaized
1 parent b67590b commit 1a79722

File tree

3 files changed

+232
-91
lines changed

3 files changed

+232
-91
lines changed

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

+22
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,17 @@
99
expand_alias_to_datasets,
1010
)
1111
from airflow.datasets.metadata import Metadata
12+
from airflow.decorators import dag, setup, task, task_group, teardown
13+
from airflow.io.path import ObjectStoragePath
14+
from airflow.io.storage import attach
15+
from airflow.models import DAG as DAGFromModel
1216
from airflow.models.baseoperator import chain, chain_linear, cross_downstream
1317
from airflow.models.baseoperatorlink import BaseOperatorLink
18+
from airflow.models.dag import DAG as DAGFromDag
1419
from airflow.timetables.datasets import DatasetOrTimeSchedule
1520
from airflow.utils.dag_parsing_context import get_parsing_context
1621

22+
# airflow
1723
DatasetFromRoot()
1824

1925
# airflow.datasets
@@ -24,6 +30,20 @@
2430
Metadata()
2531
expand_alias_to_datasets()
2632

33+
# airflow.decorators
34+
dag()
35+
task()
36+
task_group()
37+
setup()
38+
teardown()
39+
40+
# airflow.io
41+
ObjectStoragePath()
42+
attach()
43+
44+
# airflow.models
45+
DAGFromModel()
46+
2747
# airflow.models.baseoperator
2848
chain()
2949
chain_linear()
@@ -32,6 +52,8 @@
3252
# airflow.models.baseoperatolinker
3353
BaseOperatorLink()
3454

55+
# airflow.models.dag
56+
DAGFromDag()
3557
# airflow.timetables.datasets
3658
DatasetOrTimeSchedule()
3759

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

+24
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,25 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
210210
"expand_alias_to_datasets" => Replacement::Name("airflow.sdk.expand_alias_to_assets"),
211211
_ => return,
212212
},
213+
214+
// airflow.decorators
215+
["airflow", "decorators", rest @ ("dag" | "task" | "task_group" | "setup" | "teardown")] => {
216+
Replacement::SourceModuleMoved {
217+
module: "airflow.sdk",
218+
name: (*rest).to_string(),
219+
}
220+
}
221+
222+
// airflow.io
223+
["airflow", "io", "path", "ObjectStoragePath"] => Replacement::SourceModuleMoved {
224+
module: "airflow.sdk",
225+
name: "ObjectStoragePath".to_string(),
226+
},
227+
["airflow", "io", "storage", "attach"] => Replacement::SourceModuleMoved {
228+
module: "airflow.sdk.io",
229+
name: "attach".to_string(),
230+
},
231+
213232
// airflow.models.baseoperator
214233
["airflow", "models", "baseoperator", rest] => match *rest {
215234
"chain" | "chain_linear" | "cross_downstream" => Replacement::SourceModuleMoved {
@@ -221,6 +240,11 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
221240
}
222241
_ => return,
223242
},
243+
// airflow.model..DAG
244+
["airflow", "models", .., "DAG"] => Replacement::SourceModuleMoved {
245+
module: "airflow.sdk",
246+
name: "DAG".to_string(),
247+
},
224248
// airflow.timetables
225249
["airflow", "timetables", "datasets", "DatasetOrTimeSchedule"] => {
226250
Replacement::Name("airflow.timetables.assets.AssetOrTimeSchedule")

0 commit comments

Comments
 (0)