Add AIRFLOW module guard

astral-sh · Nov 26, 2024 · 7745a16 · 7745a16
1 parent 83feee3
commit 7745a16
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/crates/ruff_linter/src/rules/airflow/rules/dag_schedule_argument.rs b/crates/ruff_linter/src/rules/airflow/rules/dag_schedule_argument.rs
@@ -2,6 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::Expr;
 use ruff_python_ast::{self as ast};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -50,6 +51,10 @@ impl Violation for AirflowDagNoScheduleArgument {
 
 /// AIR301
 pub(crate) fn dag_no_schedule_argument(checker: &mut Checker, expr: &Expr) {
+    if !checker.semantic().seen_module(Modules::AIRFLOW) {
+        return;
+    }
+
     // Don't check non-call expressions.
     let Expr::Call(ast::ExprCall {
         func, arguments, ..

diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs
@@ -1289,6 +1289,7 @@ impl<'a> SemanticModel<'a> {
             "typing" => self.seen.insert(Modules::TYPING),
             "typing_extensions" => self.seen.insert(Modules::TYPING_EXTENSIONS),
             "attr" | "attrs" => self.seen.insert(Modules::ATTRS),
+            "airflow" => self.seen.insert(Modules::AIRFLOW),
             _ => {}
         }
     }
@@ -1860,6 +1861,7 @@ bitflags! {
         const FLASK = 1 << 24;
         const ATTRS = 1 << 25;
         const REGEX = 1 << 26;
+        const AIRFLOW = 1 << 27;
     }
 }