Skip to content

Commit 504c1f9

Browse files
Lee-Wmaxmynter
authored andcommitted
[airflow] Add autofix for AIR302 attribute checks (astral-sh#16977)
## Summary Add autofix logic to AIR302 check_method ## Test Plan test fixtures have been updated accordingly
1 parent 4a7a96a commit 504c1f9

File tree

2 files changed

+70
-8
lines changed

2 files changed

+70
-8
lines changed

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

+15-2
Original file line numberDiff line numberDiff line change
@@ -294,13 +294,26 @@ fn check_class_attribute(checker: &Checker, attribute_expr: &ExprAttribute) {
294294
_ => return,
295295
};
296296

297-
checker.report_diagnostic(Diagnostic::new(
297+
// Create the `Fix` first to avoid cloning `Replacement`.
298+
let fix = if let Replacement::Name(name) = replacement {
299+
Some(Fix::safe_edit(Edit::range_replacement(
300+
name.to_string(),
301+
attr.range(),
302+
)))
303+
} else {
304+
None
305+
};
306+
let mut diagnostic = Diagnostic::new(
298307
Airflow3Removal {
299308
deprecated: attr.to_string(),
300309
replacement,
301310
},
302311
attr.range(),
303-
));
312+
);
313+
if let Some(fix) = fix {
314+
diagnostic.set_fix(fix);
315+
}
316+
checker.report_diagnostic(diagnostic);
304317
}
305318

306319
/// Checks whether an Airflow 3.0–removed context key is used in a function decorated with `@task`.

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

+55-6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
---
22
source: crates/ruff_linter/src/rules/airflow/mod.rs
3-
snapshot_kind: text
43
---
54
AIR302_class_attribute.py:24:21: AIR302 `airflow.Dataset` is removed in Airflow 3.0
65
|
@@ -329,7 +328,7 @@ AIR302_class_attribute.py:50:11: AIR302 `airflow.lineage.hook.DatasetLineageInfo
329328
|
330329
= help: Use `airflow.lineage.hook.AssetLineageInfo` instead
331330

332-
AIR302_class_attribute.py:51:9: AIR302 `dataset` is removed in Airflow 3.0
331+
AIR302_class_attribute.py:51:9: AIR302 [*] `dataset` is removed in Airflow 3.0
333332
|
334333
49 | # airflow.lineage.hook
335334
50 | dl_info = DatasetLineageInfo()
@@ -340,6 +339,16 @@ AIR302_class_attribute.py:51:9: AIR302 `dataset` is removed in Airflow 3.0
340339
|
341340
= help: Use `asset` instead
342341

342+
Safe fix
343+
48 48 |
344+
49 49 | # airflow.lineage.hook
345+
50 50 | dl_info = DatasetLineageInfo()
346+
51 |-dl_info.dataset
347+
51 |+dl_info.asset
348+
52 52 |
349+
53 53 | hlc = HookLineageCollector()
350+
54 54 | hlc.create_dataset()
351+
343352
AIR302_class_attribute.py:54:5: AIR302 [*] `create_dataset` is removed in Airflow 3.0
344353
|
345354
53 | hlc = HookLineageCollector()
@@ -525,7 +534,7 @@ AIR302_class_attribute.py:79:15: AIR302 [*] `get_connections` is removed in Airf
525534
81 81 | not_an_error = NotAir302SecretError()
526535
82 82 | not_an_error.get_conn_uri()
527536

528-
AIR302_class_attribute.py:87:4: AIR302 `dataset_factories` is removed in Airflow 3.0
537+
AIR302_class_attribute.py:87:4: AIR302 [*] `dataset_factories` is removed in Airflow 3.0
529538
|
530539
85 | pm = ProvidersManager()
531540
86 | pm.initialize_providers_asset_uri_resources()
@@ -536,7 +545,17 @@ AIR302_class_attribute.py:87:4: AIR302 `dataset_factories` is removed in Airflow
536545
|
537546
= help: Use `asset_factories` instead
538547

539-
AIR302_class_attribute.py:88:4: AIR302 `dataset_factories` is removed in Airflow 3.0
548+
Safe fix
549+
84 84 | # airflow.providers_manager
550+
85 85 | pm = ProvidersManager()
551+
86 86 | pm.initialize_providers_asset_uri_resources()
552+
87 |-pm.dataset_factories
553+
87 |+pm.asset_factories
554+
88 88 | pm.dataset_factories
555+
89 89 | pm.dataset_uri_handlers
556+
90 90 | pm.dataset_to_openlineage_converters
557+
558+
AIR302_class_attribute.py:88:4: AIR302 [*] `dataset_factories` is removed in Airflow 3.0
540559
|
541560
86 | pm.initialize_providers_asset_uri_resources()
542561
87 | pm.dataset_factories
@@ -547,7 +566,17 @@ AIR302_class_attribute.py:88:4: AIR302 `dataset_factories` is removed in Airflow
547566
|
548567
= help: Use `asset_factories` instead
549568

550-
AIR302_class_attribute.py:89:4: AIR302 `dataset_uri_handlers` is removed in Airflow 3.0
569+
Safe fix
570+
85 85 | pm = ProvidersManager()
571+
86 86 | pm.initialize_providers_asset_uri_resources()
572+
87 87 | pm.dataset_factories
573+
88 |-pm.dataset_factories
574+
88 |+pm.asset_factories
575+
89 89 | pm.dataset_uri_handlers
576+
90 90 | pm.dataset_to_openlineage_converters
577+
91 91 |
578+
579+
AIR302_class_attribute.py:89:4: AIR302 [*] `dataset_uri_handlers` is removed in Airflow 3.0
551580
|
552581
87 | pm.dataset_factories
553582
88 | pm.dataset_factories
@@ -557,7 +586,17 @@ AIR302_class_attribute.py:89:4: AIR302 `dataset_uri_handlers` is removed in Airf
557586
|
558587
= help: Use `asset_uri_handlers` instead
559588

560-
AIR302_class_attribute.py:90:4: AIR302 `dataset_to_openlineage_converters` is removed in Airflow 3.0
589+
Safe fix
590+
86 86 | pm.initialize_providers_asset_uri_resources()
591+
87 87 | pm.dataset_factories
592+
88 88 | pm.dataset_factories
593+
89 |-pm.dataset_uri_handlers
594+
89 |+pm.asset_uri_handlers
595+
90 90 | pm.dataset_to_openlineage_converters
596+
91 91 |
597+
92 92 | # airflow.secrets.base_secrets
598+
599+
AIR302_class_attribute.py:90:4: AIR302 [*] `dataset_to_openlineage_converters` is removed in Airflow 3.0
561600
|
562601
88 | pm.dataset_factories
563602
89 | pm.dataset_uri_handlers
@@ -568,6 +607,16 @@ AIR302_class_attribute.py:90:4: AIR302 `dataset_to_openlineage_converters` is re
568607
|
569608
= help: Use `asset_to_openlineage_converters` instead
570609

610+
Safe fix
611+
87 87 | pm.dataset_factories
612+
88 88 | pm.dataset_factories
613+
89 89 | pm.dataset_uri_handlers
614+
90 |-pm.dataset_to_openlineage_converters
615+
90 |+pm.asset_to_openlineage_converters
616+
91 91 |
617+
92 92 | # airflow.secrets.base_secrets
618+
93 93 | base_secret_backend = BaseSecretsBackend()
619+
571620
AIR302_class_attribute.py:94:21: AIR302 [*] `get_conn_uri` is removed in Airflow 3.0
572621
|
573622
92 | # airflow.secrets.base_secrets

0 commit comments

Comments
 (0)