[airflow] Refactor AIR301 logic and fix typos (AIR301) (#17293)

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

* Simplify match conditions in AIR301
* Fix
* `airflow.datasets.manager.DatasetManager` →
`airflow.assets.manager.AssetManager`
* `airflow.www.auth.has_access_dataset` →
`airflow.www.auth.has_access_dataset`

## Test Plan

<!-- How was it tested? -->
The test fixture has been updated accordingly
This commit is contained in:
Wei Lee 2025-04-09 22:46:17 +08:00 committed by GitHub
parent 7207c86971
commit 7c81408c54
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 78 additions and 70 deletions

View file

@ -10,9 +10,13 @@ pub(crate) enum Replacement {
Name(&'static str),
Message(&'static str),
AutoImport {
path: &'static str,
module: &'static str,
name: &'static str,
},
SourceModuleMoved {
module: &'static str,
name: String,
},
}
#[derive(Debug, Eq, PartialEq)]
@ -70,10 +74,10 @@ pub(crate) fn is_guarded_by_try_except(
/// contain any [`ast::StmtImportFrom`] nodes that indicate the numpy
/// member is being imported from the non-deprecated location?
fn try_block_contains_undeprecated_import(try_node: &StmtTry, replacement: &Replacement) -> bool {
let Replacement::AutoImport { path, name } = replacement else {
let Replacement::AutoImport { module, name } = replacement else {
return false;
};
let mut import_searcher = ImportSearcher::new(path, name);
let mut import_searcher = ImportSearcher::new(module, name);
import_searcher.visit_body(&try_node.body);
import_searcher.found_import
}

View file

@ -54,16 +54,15 @@ impl Violation for Airflow3Removal {
replacement,
} = self;
match replacement {
Replacement::None => format!("`{deprecated}` is removed in Airflow 3.0"),
Replacement::Name(_) => {
Replacement::None
| Replacement::Name(_)
| Replacement::AutoImport { module: _, name: _ }
| Replacement::SourceModuleMoved { module: _, name: _ } => {
format!("`{deprecated}` is removed in Airflow 3.0")
}
Replacement::Message(message) => {
format!("`{deprecated}` is removed in Airflow 3.0; {message}")
}
Replacement::AutoImport { path: _, name: _ } => {
format!("`{deprecated}` is removed in Airflow 3.0")
}
}
}
@ -71,7 +70,12 @@ impl Violation for Airflow3Removal {
let Airflow3Removal { replacement, .. } = self;
match replacement {
Replacement::Name(name) => Some(format!("Use `{name}` instead")),
Replacement::AutoImport { path, name } => Some(format!("Use `{path}.{name}` instead")),
Replacement::AutoImport { module, name } => {
Some(format!("Use `{module}.{name}` instead"))
}
Replacement::SourceModuleMoved { module, name } => {
Some(format!("Use `{module}.{name}` instead"))
}
_ => None,
}
}
@ -575,7 +579,7 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
}
["airflow", "api_connexion", "security", "requires_access_dataset"] => {
Replacement::AutoImport {
path: "airflow.api_connexion.security",
module: "airflow.api_connexion.security",
name: "requires_access_asset",
}
}
@ -593,16 +597,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
}
// airflow.configuration
["airflow", "configuration", rest @ ..] => match &rest {
["get"] => Replacement::Name("airflow.configuration.conf.get"),
["getboolean"] => Replacement::Name("airflow.configuration.conf.getboolean"),
["getfloat"] => Replacement::Name("airflow.configuration.conf.getfloat"),
["getint"] => Replacement::Name("airflow.configuration.conf.getint"),
["has_option"] => Replacement::Name("airflow.configuration.conf.has_option"),
["remove_option"] => Replacement::Name("airflow.configuration.conf.remove_option"),
["as_dict"] => Replacement::Name("airflow.configuration.conf.as_dict"),
["set"] => Replacement::Name("airflow.configuration.conf.set"),
_ => return,
["airflow", "configuration", rest @ ("as_dict" | "get" | "getboolean" | "getfloat" | "getint" | "has_option"
| "remove_option" | "set")] => Replacement::SourceModuleMoved {
module: "airflow.configuration.conf",
name: (*rest).to_string(),
},
// airflow.contrib.*
@ -610,26 +608,27 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
Replacement::Message("The whole `airflow.contrib` module has been removed.")
}
// airflow.datasets.manager
["airflow", "datasets", "manager", rest] => match *rest {
"DatasetManager" => Replacement::Name("airflow.assets.manager.AssetManager"),
"dataset_manager" => Replacement::Name("airflow.assets.manager.asset_manager"),
"resolve_dataset_manager" => Replacement::Name("airflow.assets.resolve_asset_manager"),
_ => return,
},
["airflow", "datasets", "metadata", "Metadata"] => {
Replacement::Name("airflow.sdk.Metadata")
}
// airflow.datasets
["airflow", "Dataset"] | ["airflow", "datasets", "Dataset"] => Replacement::AutoImport {
path: "airflow.sdk",
module: "airflow.sdk",
name: "Asset",
},
["airflow", "datasets", rest @ ..] => match &rest {
["DatasetAliasEvent"] => Replacement::None,
["DatasetAlias"] => Replacement::Name("airflow.sdk.AssetAlias"),
["DatasetAll"] => Replacement::Name("airflow.sdk.AssetAll"),
["DatasetAny"] => Replacement::Name("airflow.sdk.AssetAny"),
["expand_alias_to_datasets"] => Replacement::Name("airflow.sdk.expand_alias_to_assets"),
["metadata", "Metadata"] => Replacement::Name("airflow.sdk.Metadata"),
// airflow.datasets.manager
["manager", "DatasetManager"] => Replacement::Name("airflow.assets.AssetManager"),
["manager", "dataset_manager"] => {
Replacement::Name("airflow.assets.manager.asset_manager")
}
["manager", "resolve_dataset_manager"] => {
Replacement::Name("airflow.assets.resolve_asset_manager")
}
["airflow", "datasets", rest] => match *rest {
"DatasetAliasEvent" => Replacement::None,
"DatasetAlias" => Replacement::Name("airflow.sdk.AssetAlias"),
"DatasetAll" => Replacement::Name("airflow.sdk.AssetAll"),
"DatasetAny" => Replacement::Name("airflow.sdk.AssetAny"),
"expand_alias_to_datasets" => Replacement::Name("airflow.sdk.expand_alias_to_assets"),
_ => return,
},
@ -644,38 +643,39 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
}
// airflow.listeners.spec
["airflow", "listeners", "spec", "dataset", rest @ ..] => match &rest {
["on_dataset_created"] => {
// TODO: this is removed
["airflow", "listeners", "spec", "dataset", rest] => match *rest {
"on_dataset_created" => {
Replacement::Name("airflow.listeners.spec.asset.on_asset_created")
}
["on_dataset_changed"] => {
"on_dataset_changed" => {
Replacement::Name("airflow.listeners.spec.asset.on_asset_changed")
}
_ => return,
},
// airflow.metrics.validators
["airflow", "metrics", "validators", rest @ ..] => match &rest {
["AllowListValidator"] => {
["airflow", "metrics", "validators", rest] => match *rest {
"AllowListValidator" => {
Replacement::Name("airflow.metrics.validators.PatternAllowListValidator")
}
["BlockListValidator"] => {
"BlockListValidator" => {
Replacement::Name("airflow.metrics.validators.PatternBlockListValidator")
}
_ => return,
},
// airflow.models.baseoperator
["airflow", "models", "baseoperator", "chain"] => Replacement::Name("airflow.sdk.chain"),
["airflow", "models", "baseoperator", "chain_linear"] => {
Replacement::Name("airflow.sdk.chain_linear")
}
["airflow", "models", "baseoperator", "cross_downstream"] => {
Replacement::Name("airflow.sdk.cross_downstream")
}
["airflow", "models", "baseoperatorlink", "BaseOperatorLink"] => {
Replacement::Name("airflow.sdk.definitions.baseoperatorlink.BaseOperatorLink")
}
["airflow", "models", "baseoperator", rest] => match *rest {
"chain" | "chain_linear" | "cross_downstream" => Replacement::SourceModuleMoved {
module: "airflow.sdk",
name: (*rest).to_string(),
},
"BaseOperatorLink" => {
Replacement::Name("airflow.sdk.definitions.baseoperatorlink.BaseOperatorLink")
}
_ => return,
},
// airflow.notifications
["airflow", "notifications", "basenotifier", "BaseNotifier"] => {
@ -765,11 +765,12 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
},
// airflow.www
// TODO: www has been removed
["airflow", "www", "auth", "has_access"] => {
Replacement::Name("airflow.www.auth.has_access_*")
}
["airflow", "www", "auth", "has_access_dataset"] => {
Replacement::Name("airflow.www.auth.has_access_dataset.has_access_asset")
Replacement::Name("airflow.www.auth.has_access_dataset")
}
["airflow", "www", "utils", "get_sensitive_variables_fields"] => {
Replacement::Name("airflow.utils.log.secrets_masker.get_sensitive_variables_fields")
@ -796,14 +797,15 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
},
// airflow.providers.common.io
["airflow", "providers", "common", "io", rest @ ..] => match &rest {
["datasets", "file", "create_dataset"] => {
// airflow.providers.common.io.datasets.file
["airflow", "providers", "common", "io", "datasets", "file", rest] => match *rest {
"create_dataset" => {
Replacement::Name("airflow.providers.common.io.assets.file.create_asset")
}
["datasets", "file", "convert_dataset_to_openlineage"] => Replacement::Name(
"convert_dataset_to_openlineage" => Replacement::Name(
"airflow.providers.common.io.assets.file.convert_asset_to_openlineage",
),
["datasets", "file", "sanitize_uri"] => {
"sanitize_uri" => {
Replacement::Name("airflow.providers.common.io.assets.file.sanitize_uri")
}
_ => return,
@ -817,17 +819,18 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
}
// airflow.providers.google
["airflow", "providers", "google", rest @ ..] => match &rest {
["datasets", "bigquery", "create_dataset"] => {
// airflow.providers.google.datasets
["airflow", "providers", "google", "datasets", rest @ ..] => match &rest {
["bigquery", "create_dataset"] => {
Replacement::Name("airflow.providers.google.assets.bigquery.create_asset")
}
["datasets", "gcs", "create_dataset"] => {
["gcs", "create_dataset"] => {
Replacement::Name("airflow.providers.google.assets.gcs.create_asset")
}
["datasets", "gcs", "convert_dataset_to_openlineage"] => Replacement::Name(
["gcs", "convert_dataset_to_openlineage"] => Replacement::Name(
"airflow.providers.google.assets.gcs.convert_asset_to_openlineage",
),
["datasets", "gcs", "sanitize_uri"] => {
["gcs", "sanitize_uri"] => {
Replacement::Name("airflow.providers.google.assets.gcs.sanitize_uri")
}
_ => return,
@ -844,12 +847,13 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
}
// airflow.providers.openlineage
["airflow", "providers", "openlineage", rest @ ..] => match &rest {
["utils", "utils", "DatasetInfo"] => {
// airflow.providers.openlineage.utils.utils
["airflow", "providers", "openlineage", "utils", "utils", rest] => match *rest {
"DatasetInfo" => {
Replacement::Name("airflow.providers.openlineage.utils.utils.AssetInfo")
}
["utils", "utils", "translate_airflow_dataset"] => Replacement::Name(
"translate_airflow_dataset" => Replacement::Name(
"airflow.providers.openlineage.utils.utils.translate_airflow_asset",
),
_ => return,
@ -875,10 +879,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
range,
);
if let Replacement::AutoImport { path, name } = replacement {
if let Replacement::AutoImport { module, name } = replacement {
diagnostic.try_set_fix(|| {
let (import_edit, binding) = checker.importer().get_or_import_symbol(
&ImportRequest::import_from(path, name),
&ImportRequest::import_from(module, name),
expr.start(),
checker.semantic(),
)?;

View file

@ -244,7 +244,7 @@ AIR301_class_attribute.py:42:6: AIR301 `airflow.datasets.manager.DatasetManager`
43 | dm.register_dataset_change()
44 | dm.create_datasets()
|
= help: Use `airflow.assets.AssetManager` instead
= help: Use `airflow.assets.manager.AssetManager` instead
AIR301_class_attribute.py:43:4: AIR301 [*] `register_dataset_change` is removed in Airflow 3.0
|

View file

@ -321,7 +321,7 @@ AIR301_names.py:118:1: AIR301 `airflow.datasets.manager.DatasetManager` is remov
119 | dataset_manager
120 | resolve_dataset_manager
|
= help: Use `airflow.assets.AssetManager` instead
= help: Use `airflow.assets.manager.AssetManager` instead
AIR301_names.py:119:1: AIR301 `airflow.datasets.manager.dataset_manager` is removed in Airflow 3.0
|
@ -876,7 +876,7 @@ AIR301_names.py:264:1: AIR301 `airflow.www.auth.has_access_dataset` is removed i
265 |
266 | # airflow.www.utils
|
= help: Use `airflow.www.auth.has_access_dataset.has_access_asset` instead
= help: Use `airflow.www.auth.has_access_dataset` instead
AIR301_names.py:267:1: AIR301 `airflow.www.utils.get_sensitive_variables_fields` is removed in Airflow 3.0
|