Add AIR001: task variable name should be same as task_id arg

crates/ruff/src/checkers/ast/mod.rs

@@ -42,7 +42,7 @@ use crate::importer::Importer;
 use crate::noqa::NoqaMapping;
 use crate::registry::{AsRule, Rule};
 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,
@@ -1799,6 +1799,14 @@ where
                     }
                 }
 
+                if self.settings.rules.enabled(Rule::TaskVariableNameNotTaskId) {
+                    if let Some(diagnostic) =
+                        airflow::rules::task_variable_name(self, targets, value)
+                    {
+                        self.diagnostics.push(diagnostic);
+                    }
+                }
+
                 if self.is_stub {
                     if self.settings.rules.any_enabled(&[
                         Rule::UnprefixedTypeParam,

crates/ruff/src/codes.rs

@@ -745,6 +745,9 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Flake8Todo, "006") => (RuleGroup::Unspecified, Rule::InvalidTodoCapitalization),
         (Flake8Todo, "007") => (RuleGroup::Unspecified, Rule::MissingSpaceAfterTodoColon),
 
+        // airflow
+        (Airflow, "001") => (RuleGroup::Unspecified, Rule::TaskVariableNameNotTaskId),
+
         _ => return None,
     })
 }

crates/ruff/src/registry.rs

@@ -664,6 +664,8 @@ ruff_macros::register_rules!(
     rules::flake8_todos::rules::MissingTodoDescription,
     rules::flake8_todos::rules::InvalidTodoCapitalization,
     rules::flake8_todos::rules::MissingSpaceAfterTodoColon,
+    // airflow
+    rules::airflow::rules::TaskVariableNameNotTaskId,
 );
 
 pub trait AsRule {
@@ -834,6 +836,9 @@ pub enum Linter {
     /// NumPy-specific rules
     #[prefix = "NPY"]
     Numpy,
+    /// [Airflow](https://pypi.org/project/apache-airflow/)
+    #[prefix = "AIR"]
+    Airflow,
     /// Ruff-specific rules
     #[prefix = "RUF"]
     Ruff,

crates/ruff/src/rules/airflow/mod.rs

@@ -0,0 +1,69 @@
+//! Airflow-specific rules
+pub(crate) mod rules;
+
+#[cfg(test)]
+mod tests {
+    use std::path::Path;
+
+    use rustpython_parser::lexer::LexResult;
+    use test_case::test_case;
+    use textwrap::dedent;
+
+    use ruff_python_ast::source_code::{Indexer, Locator, Stylist};
+
+    use crate::linter::{check_path, LinterResult};
+    use crate::registry::{AsRule, Linter, Rule};
+    use crate::settings::flags;
+    use crate::{directives, settings};
+
+    fn rule_code(contents: &str, expected: &[Rule]) {
+        let contents = dedent(contents);
+        let settings = settings::Settings::for_rules(&Linter::Airflow);
+        let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&contents);
+        let locator = Locator::new(&contents);
+        let stylist = Stylist::from_tokens(&tokens, &locator);
+        let indexer = Indexer::from_tokens(&tokens, &locator);
+        let directives = directives::extract_directives(
+            &tokens,
+            directives::Flags::from_settings(&settings),
+            &locator,
+            &indexer,
+        );
+        let LinterResult {
+            data: (diagnostics, _imports),
+            ..
+        } = check_path(
+            Path::new("<filename>"),
+            None,
+            tokens,
+            &locator,
+            &stylist,
+            &indexer,
+            &directives,
+            &settings,
+            flags::Noqa::Enabled,
+        );
+        let actual: Vec<Rule> = diagnostics
+            .into_iter()
+            .map(|diagnostic| diagnostic.kind.rule())
+            .collect();
+        assert_eq!(actual, expected);
+    }
+
+    #[test_case(r#"
+        from airflow.operators import PythonOperator
+        my_task = PythonOperator(task_id="my_task")
+    "#, &[]; "AIR001_pass")]
+    #[test_case(r#"
+        from airflow.operators import PythonOperator
+        incorrect_name = PythonOperator(task_id="my_task")
+    "#, &[Rule::TaskVariableNameNotTaskId]; "AIR001_fail")]
+    #[test_case(r#"
+        from my_module import MyClass
+        incorrect_name = MyClass(task_id="my_task")
+        "#, &[]; "AIR001_noop")]
+
+    fn test_airflow(code: &str, expected: &[Rule]) {
+        rule_code(code, expected);
+    }
+}

crates/ruff/src/rules/airflow/rules/mod.rs

@@ -0,0 +1,3 @@
+mod tasks;
+
+pub(crate) use tasks::{task_variable_name, TaskVariableNameNotTaskId};

crates/ruff/src/rules/airflow/rules/tasks.rs

@@ -0,0 +1,100 @@
+use rustpython_parser::ast;
+use rustpython_parser::ast::{Expr, Ranged};
+
+use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_macros::{derive_message_formats, violation};
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks that the task variable name is the same as the task_id.
+///
+/// ## Why is this bad?
+/// For consistency, the task variable name should be the same as the task_id.
+/// This makes it easier for you and others to understand the code.
+///
+/// ## Example
+/// ```python
+/// incorrect_name = SomeOperator(task_id="my_task")
+/// async def fetch():
+/// ```
+///
+/// Use instead:
+/// ```python
+/// my_task = SomeOperator(task_id="my_task")
+/// ```
+#[violation]
+pub struct TaskVariableNameNotTaskId;
+
+impl Violation for TaskVariableNameNotTaskId {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Task variable name should be the same as the task_id")
+    }
+}
+
+/// AIR001
+pub(crate) fn task_variable_name(
+    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 call = match value {
+        Expr::Call(call) => call,
+        _ => return None,
+    };
+
+    // if the function doesn't come from airflow, we can't do anything
+    let func_name = match call.func.as_name_expr() {
+        Some(name) => name.id.as_str(),
+        None => return None,
+    };
+    let fully_qualified_func_name = match checker.ctx.find_binding(func_name) {
+        Some(call_path) => match call_path.kind.as_from_importation() {
+            Some(from_importation) => &from_importation.full_name,
+            None => return None,
+        },
+        None => return None,
+    };
+
+    if !fully_qualified_func_name.starts_with("airflow.") {
+        return None;
+    }
+
+    // if the call doesn't have a task_id, don't do anything
+    let task_id_arg = match call.keywords.iter().find(|keyword| match &keyword.arg {
+        Some(arg_name) => arg_name == "task_id",
+        _ => false,
+    }) {
+        Some(keyword) => keyword,
+        None => return None,
+    };
+
+    // get the task_id value
+    let task_id_arg_value = match &task_id_arg.value {
+        Expr::Constant(constant) => match constant.value.as_str() {
+            Some(string) => string,
+            None => return None,
+        },
+        _ => return None,
+    };
+
+    // if the target name is the same as the task_id, no violation
+    if &id.to_string() == task_id_arg_value {
+        return None;
+    }
+
+    // if the target name is not the same as the task_id, violation
+    Some(Diagnostic::new(TaskVariableNameNotTaskId, target.range()))
+}

crates/ruff/src/rules/mod.rs

@@ -1,4 +1,5 @@
 #![allow(clippy::useless_format)]
+pub mod airflow;
 pub mod eradicate;
 pub mod flake8_2020;
 pub mod flake8_annotations;