From 957b3e0c9f37414bd68fa4e45f70fcbabe8637e5 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Wed, 27 Nov 2024 12:49:16 +0530 Subject: [PATCH] Check that `airflow` module is seen for `AIR001` --- .../src/rules/airflow/rules/task_variable_name.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/ruff_linter/src/rules/airflow/rules/task_variable_name.rs b/crates/ruff_linter/src/rules/airflow/rules/task_variable_name.rs index 9836f230e7b9ba..097a8da2c89463 100644 --- a/crates/ruff_linter/src/rules/airflow/rules/task_variable_name.rs +++ b/crates/ruff_linter/src/rules/airflow/rules/task_variable_name.rs @@ -2,6 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation}; use ruff_macros::{derive_message_formats, violation}; use ruff_python_ast as ast; use ruff_python_ast::Expr; +use ruff_python_semantic::Modules; use ruff_text_size::Ranged; use crate::checkers::ast::Checker; @@ -49,6 +50,10 @@ pub(crate) fn variable_name_task_id( targets: &[Expr], value: &Expr, ) -> Option { + if !checker.semantic().seen_module(Modules::AIRFLOW) { + return None; + } + // If we have more than one target, we can't do anything. let [target] = targets else { return None; @@ -69,7 +74,7 @@ pub(crate) fn variable_name_task_id( if !checker .semantic() .resolve_qualified_name(func) - .is_some_and(|qualified_name| matches!(qualified_name.segments()[0], "airflow")) + .is_some_and(|qualified_name| matches!(qualified_name.segments(), ["airflow", ..])) { return None; }