From 5d7c17c20adbbea270234fe65b4849fc76a786d6 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Wed, 3 Sep 2025 22:12:11 +0800
Subject: [PATCH] [`airflow`] Convert `DatasetOrTimeSchedule(datasets=...)` to
 `AssetOrTimeSchedule(assets=...)` (`AIR311`) (#20202)

<!--
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? -->

update the argument `datasets` as `assets`

## Test Plan

<!-- How was it tested? -->

update fixture accordingly
---
 .../test/fixtures/airflow/AIR311_names.py     |  2 +-
 .../airflow/rules/suggested_to_update_3_0.rs  |  3 +++
 ...irflow__tests__AIR311_AIR311_names.py.snap | 27 ++++++++++++++++---
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR311_names.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR311_names.py
index 9b4ce1e38e..2b415aa407 100644
--- a/crates/ruff_linter/resources/test/fixtures/airflow/AIR311_names.py
+++ b/crates/ruff_linter/resources/test/fixtures/airflow/AIR311_names.py
@@ -70,7 +70,7 @@ from airflow.timetables.datasets import DatasetOrTimeSchedule
 from airflow.utils.dag_parsing_context import get_parsing_context
 
 # airflow.timetables.datasets
-DatasetOrTimeSchedule()
+DatasetOrTimeSchedule(datasets=[])
 
 # airflow.utils.dag_parsing_context
 get_parsing_context()
diff --git a/crates/ruff_linter/src/rules/airflow/rules/suggested_to_update_3_0.rs b/crates/ruff_linter/src/rules/airflow/rules/suggested_to_update_3_0.rs
index 41bc35d352..a7c83abbca 100644
--- a/crates/ruff_linter/src/rules/airflow/rules/suggested_to_update_3_0.rs
+++ b/crates/ruff_linter/src/rules/airflow/rules/suggested_to_update_3_0.rs
@@ -157,6 +157,9 @@ fn check_call_arguments(checker: &Checker, qualified_name: &QualifiedName, argum
         ["airflow", .., "DAG" | "dag"] => {
             diagnostic_for_argument(checker, arguments, "sla_miss_callback", None);
         }
+        ["airflow", "timetables", "datasets", "DatasetOrTimeSchedule"] => {
+            diagnostic_for_argument(checker, arguments, "datasets", Some("assets"));
+        }
         segments => {
             if is_airflow_builtin_or_provider(segments, "operators", "Operator") {
                 diagnostic_for_argument(checker, arguments, "sla", None);
diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR311_AIR311_names.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR311_AIR311_names.py.snap
index 3fcbf77530..a3bced5174 100644
--- a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR311_AIR311_names.py.snap
+++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR311_AIR311_names.py.snap
@@ -558,7 +558,7 @@ AIR311 [*] `airflow.timetables.datasets.DatasetOrTimeSchedule` is removed in Air
   --> AIR311_names.py:73:1
    |
 72 | # airflow.timetables.datasets
-73 | DatasetOrTimeSchedule()
+73 | DatasetOrTimeSchedule(datasets=[])
    | ^^^^^^^^^^^^^^^^^^^^^
 74 |
 75 | # airflow.utils.dag_parsing_context
@@ -570,12 +570,31 @@ help: Use `AssetOrTimeSchedule` from `airflow.timetables.assets` instead.
 71 + from airflow.timetables.assets import AssetOrTimeSchedule
 72 | 
 73 | # airflow.timetables.datasets
-   - DatasetOrTimeSchedule()
-74 + AssetOrTimeSchedule()
+   - DatasetOrTimeSchedule(datasets=[])
+74 + AssetOrTimeSchedule(datasets=[])
 75 | 
 76 | # airflow.utils.dag_parsing_context
 77 | get_parsing_context()
 
+AIR311 [*] `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:73:23
+   |
+72 | # airflow.timetables.datasets
+73 | DatasetOrTimeSchedule(datasets=[])
+   |                       ^^^^^^^^
+74 |
+75 | # airflow.utils.dag_parsing_context
+   |
+help: Use `assets` instead
+70 | from airflow.utils.dag_parsing_context import get_parsing_context
+71 | 
+72 | # airflow.timetables.datasets
+   - DatasetOrTimeSchedule(datasets=[])
+73 + DatasetOrTimeSchedule(assets=[])
+74 | 
+75 | # airflow.utils.dag_parsing_context
+76 | get_parsing_context()
+
 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.
   --> AIR311_names.py:76:1
    |
@@ -593,7 +612,7 @@ help: Use `get_parsing_context` from `airflow.sdk` instead.
 70 + from airflow.sdk import get_parsing_context
 71 | 
 72 | # airflow.timetables.datasets
-73 | DatasetOrTimeSchedule()
+73 | DatasetOrTimeSchedule(datasets=[])
 note: This is an unsafe fix and may change runtime behavior
 
 AIR311 [*] `airflow.decorators.base.DecoratedMappedOperator` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.