-
Notifications
You must be signed in to change notification settings - Fork 888
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add AIR001: task variable name should be same as task_id arg #4687
Changes from all commits
ff812fe
541d798
560ff88
f8481f4
fa93c70
67adf70
8708257
7776aca
f62a2f2
403fb69
aed21fb
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,16 @@ | ||
from airflow.operators import PythonOperator | ||
|
||
|
||
def my_callable(): | ||
pass | ||
|
||
|
||
my_task = PythonOperator(task_id="my_task", callable=my_callable) | ||
my_task_2 = PythonOperator(callable=my_callable, task_id="my_task_2") | ||
|
||
incorrect_name = PythonOperator(task_id="my_task") | ||
incorrect_name_2 = PythonOperator(callable=my_callable, task_id="my_task_2") | ||
|
||
from my_module import MyClass | ||
|
||
incorrect_name = MyClass(task_id="my_task") |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -45,7 +45,7 @@ use crate::noqa::NoqaMapping; | |
use crate::registry::{AsRule, Rule}; | ||
use crate::rules::flake8_builtins::helpers::AnyShadowing; | ||
use crate::rules::{ | ||
flake8_2020, flake8_annotations, flake8_async, flake8_bandit, flake8_blind_except, | ||
airflow, flake8_2020, flake8_annotations, flake8_async, flake8_bandit, flake8_blind_except, | ||
flake8_boolean_trap, flake8_bugbear, flake8_builtins, flake8_comprehensions, flake8_datetimez, | ||
flake8_debugger, flake8_django, flake8_errmsg, flake8_future_annotations, flake8_gettext, | ||
flake8_implicit_str_concat, flake8_import_conventions, flake8_logging_format, flake8_pie, | ||
|
@@ -1633,27 +1633,23 @@ where | |
pycodestyle::rules::lambda_assignment(self, target, value, None, stmt); | ||
} | ||
} | ||
|
||
if self.enabled(Rule::AssignmentToOsEnviron) { | ||
flake8_bugbear::rules::assignment_to_os_environ(self, targets); | ||
} | ||
|
||
if self.enabled(Rule::HardcodedPasswordString) { | ||
if let Some(diagnostic) = | ||
flake8_bandit::rules::assign_hardcoded_password_string(value, targets) | ||
{ | ||
self.diagnostics.push(diagnostic); | ||
} | ||
} | ||
|
||
if self.enabled(Rule::GlobalStatement) { | ||
for target in targets.iter() { | ||
if let Expr::Name(ast::ExprName { id, .. }) = target { | ||
pylint::rules::global_statement(self, id); | ||
} | ||
} | ||
} | ||
|
||
if self.enabled(Rule::UselessMetaclassType) { | ||
pyupgrade::rules::useless_metaclass_type(self, stmt, value, targets); | ||
} | ||
|
@@ -1670,13 +1666,22 @@ where | |
if self.enabled(Rule::UnpackedListComprehension) { | ||
pyupgrade::rules::unpacked_list_comprehension(self, targets, value); | ||
} | ||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. (Unrelated, but the newlines here felt inconsistent, so just removed.) |
||
if self.enabled(Rule::PandasDfVariableName) { | ||
if let Some(diagnostic) = pandas_vet::rules::assignment_to_df(targets) { | ||
self.diagnostics.push(diagnostic); | ||
} | ||
} | ||
|
||
if self | ||
.settings | ||
.rules | ||
.enabled(Rule::AirflowVariableNameTaskIdMismatch) | ||
{ | ||
if let Some(diagnostic) = | ||
airflow::rules::variable_name_task_id(self, targets, value) | ||
{ | ||
self.diagnostics.push(diagnostic); | ||
} | ||
} | ||
if self.is_stub { | ||
if self.any_enabled(&[ | ||
Rule::UnprefixedTypeParam, | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
//! Airflow-specific rules. | ||
pub(crate) mod rules; | ||
|
||
#[cfg(test)] | ||
mod tests { | ||
use std::path::Path; | ||
|
||
use anyhow::Result; | ||
use test_case::test_case; | ||
|
||
use crate::registry::Rule; | ||
use crate::test::test_path; | ||
use crate::{assert_messages, settings}; | ||
|
||
#[test_case(Rule::AirflowVariableNameTaskIdMismatch, Path::new("AIR001.py"); "AIR001")] | ||
fn rules(rule_code: Rule, path: &Path) -> Result<()> { | ||
let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy()); | ||
let diagnostics = test_path( | ||
Path::new("airflow").join(path).as_path(), | ||
&settings::Settings::for_rule(rule_code), | ||
)?; | ||
assert_messages!(snapshot, diagnostics); | ||
Ok(()) | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,3 @@ | ||
mod task_variable_name; | ||
|
||
pub(crate) use task_variable_name::{variable_name_task_id, AirflowVariableNameTaskIdMismatch}; |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,102 @@ | ||
use rustpython_parser::ast; | ||
use rustpython_parser::ast::{Expr, Ranged}; | ||
|
||
use ruff_diagnostics::{Diagnostic, Violation}; | ||
use ruff_macros::{derive_message_formats, violation}; | ||
use ruff_python_ast::prelude::Constant; | ||
|
||
use crate::checkers::ast::Checker; | ||
|
||
/// ## What it does | ||
/// Checks that the task variable name matches the `task_id` value for | ||
/// Airflow Operators. | ||
/// | ||
/// ## Why is this bad? | ||
/// When initializing an Airflow Operator, for consistency, the variable | ||
/// name should match the `task_id` value. This makes it easier to | ||
/// follow the flow of the DAG. | ||
/// | ||
/// ## Example | ||
/// ```python | ||
/// from airflow.operators import PythonOperator | ||
/// | ||
/// | ||
/// incorrect_name = PythonOperator(task_id="my_task") | ||
/// ``` | ||
/// | ||
/// Use instead: | ||
/// ```python | ||
/// from airflow.operators import PythonOperator | ||
/// | ||
/// | ||
/// my_task = PythonOperator(task_id="my_task") | ||
/// ``` | ||
#[violation] | ||
pub struct AirflowVariableNameTaskIdMismatch { | ||
task_id: String, | ||
} | ||
|
||
impl Violation for AirflowVariableNameTaskIdMismatch { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I renamed this to prefix with "Airflow" to match the "Django" and "NumPy" conventions. |
||
#[derive_message_formats] | ||
fn message(&self) -> String { | ||
let AirflowVariableNameTaskIdMismatch { task_id } = self; | ||
format!("Task variable name should match the `task_id`: \"{task_id}\"") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. (Added the |
||
} | ||
} | ||
|
||
/// AIR001 | ||
pub(crate) fn variable_name_task_id( | ||
checker: &mut Checker, | ||
targets: &[Expr], | ||
value: &Expr, | ||
) -> Option<Diagnostic> { | ||
// If we have more than one target, we can't do anything. | ||
if targets.len() != 1 { | ||
return None; | ||
} | ||
|
||
let target = &targets[0]; | ||
let Expr::Name(ast::ExprName { id, .. }) = target else { | ||
return None; | ||
}; | ||
|
||
// If the value is not a call, we can't do anything. | ||
let Expr::Call(ast::ExprCall { func, keywords, .. }) = value else { | ||
return None; | ||
}; | ||
|
||
// If the function doesn't come from Airflow, we can't do anything. | ||
if !checker | ||
.semantic_model() | ||
.resolve_call_path(func) | ||
.map_or(false, |call_path| matches!(call_path[0], "airflow")) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Modified this to use import airflow
foo = airflow.PythonOperator(...)
import airflow as Foo
foo = Foo.PythonOperator(...) |
||
{ | ||
return None; | ||
} | ||
|
||
// If the call doesn't have a `task_id` keyword argument, we can't do anything. | ||
let keyword = keywords | ||
.iter() | ||
.find(|keyword| keyword.arg.as_ref().map_or(false, |arg| arg == "task_id"))?; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can use |
||
|
||
// If the keyword argument is not a string, we can't do anything. | ||
let task_id = match &keyword.value { | ||
Expr::Constant(constant) => match &constant.value { | ||
Constant::Str(value) => value, | ||
_ => return None, | ||
}, | ||
_ => return None, | ||
}; | ||
|
||
// If the target name is the same as the task_id, no violation. | ||
if id == task_id { | ||
return None; | ||
} | ||
|
||
Some(Diagnostic::new( | ||
AirflowVariableNameTaskIdMismatch { | ||
task_id: task_id.to_string(), | ||
}, | ||
target.range(), | ||
)) | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,22 @@ | ||
--- | ||
source: crates/ruff/src/rules/airflow/mod.rs | ||
--- | ||
AIR001.py:11:1: AIR001 Task variable name should match the `task_id`: "my_task" | ||
| | ||
11 | my_task_2 = PythonOperator(callable=my_callable, task_id="my_task_2") | ||
12 | | ||
13 | incorrect_name = PythonOperator(task_id="my_task") | ||
| ^^^^^^^^^^^^^^ AIR001 | ||
14 | incorrect_name_2 = PythonOperator(callable=my_callable, task_id="my_task_2") | ||
| | ||
|
||
AIR001.py:12:1: AIR001 Task variable name should match the `task_id`: "my_task_2" | ||
| | ||
12 | incorrect_name = PythonOperator(task_id="my_task") | ||
13 | incorrect_name_2 = PythonOperator(callable=my_callable, task_id="my_task_2") | ||
| ^^^^^^^^^^^^^^^^ AIR001 | ||
14 | | ||
15 | from my_module import MyClass | ||
| | ||
|
||
|
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I decided to link to the Pylint plugin, since this list refers to tools we've reimplemented, and it felt more appropriate than "Airflow" -- wdyt?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
yeah this makes sense to me! hopefully we get to a point where we're linting for more than just pylint-airflow though 🙂