Track top-level module imports in the semantic model

astral-sh · Feb 2, 2024 · e09bf0a · e09bf0a
1 parent ea1c089
commit e09bf0a
Show file tree

Hide file tree

Showing 55 changed files with 394 additions and 99 deletions.
diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs
@@ -319,7 +319,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 numpy::rules::numpy_2_0_deprecation(checker, expr);
             }
             if checker.enabled(Rule::DeprecatedMockImport) {
-                pyupgrade::rules::deprecated_mock_attribute(checker, expr);
+                pyupgrade::rules::deprecated_mock_attribute(checker, attribute);
             }
             if checker.enabled(Rule::SixPY3) {
                 flake8_2020::rules::name_or_attribute(checker, expr);
@@ -336,7 +336,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UndocumentedWarn) {
                 flake8_logging::rules::undocumented_warn(checker, expr);
             }
-            pandas_vet::rules::attr(checker, attribute);
+            if checker.enabled(Rule::PandasUseOfDotValues) {
+                pandas_vet::rules::attr(checker, attribute);
+            }
         }
         Expr::Call(
             call @ ast::ExprCall {
@@ -976,7 +978,6 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnsortedDunderAll) {
                 ruff::rules::sort_dunder_all_extend_call(checker, call);
             }
-
             if checker.enabled(Rule::DefaultFactoryKwarg) {
                 ruff::rules::default_factory_kwarg(checker, call);
             }

diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs
@@ -349,13 +349,6 @@ where
             }
         }
 
-        // Track each top-level import, to guide import insertions.
-        if matches!(stmt, Stmt::Import(_) | Stmt::ImportFrom(_)) {
-            if self.semantic.at_top_level() {
-                self.importer.visit_import(stmt);
-            }
-        }
-
         // Store the flags prior to any further descent, so that we can restore them after visiting
         // the node.
         let flags_snapshot = self.semantic.flags;
@@ -371,14 +364,22 @@ where
                 self.handle_node_load(target);
             }
             Stmt::Import(ast::StmtImport { names, range: _ }) => {
+                if self.semantic.at_top_level() {
+                    self.importer.visit_import(stmt);
+                }
+
                 for alias in names {
-                    if alias.name.contains('.') && alias.asname.is_none() {
-                        // Given `import foo.bar`, `name` would be "foo", and `qualified_name` would be
-                        // "foo.bar".
-                        let name = alias.name.split('.').next().unwrap();
+                    // Given `import foo.bar`, `module` would be "foo", and `call_path` would be
+                    // `["foo", "bar"]`.
+                    let module = alias.name.split('.').next().unwrap();
+
+                    // Mark the top-level module as "seen" by the semantic model.
+                    self.semantic.add_module(module);
+
+                    if alias.asname.is_none() && alias.name.contains('.') {
                         let call_path: Box<[&str]> = alias.name.split('.').collect();
                         self.add_binding(
-                            name,
+                            module,
                             alias.identifier(),
                             BindingKind::SubmoduleImport(SubmoduleImport { call_path }),
                             BindingFlags::EXTERNAL,
@@ -413,8 +414,20 @@ where
                 level,
                 range: _,
             }) => {
+                if self.semantic.at_top_level() {
+                    self.importer.visit_import(stmt);
+                }
+
                 let module = module.as_deref();
                 let level = *level;
+
+                // Mark the top-level module as "seen" by the semantic model.
+                if level.map_or(true, |level| level == 0) {
+                    if let Some(module) = module.and_then(|module| module.split('.').next()) {
+                        self.semantic.add_module(module);
+                    }
+                }
+
                 for alias in names {
                     if let Some("__future__") = module {
                         let name = alias.asname.as_ref().unwrap_or(&alias.name);

diff --git a/crates/ruff_linter/src/rules/flake8_2020/rules/name_or_attribute.rs b/crates/ruff_linter/src/rules/flake8_2020/rules/name_or_attribute.rs
@@ -1,7 +1,7 @@
-use ruff_python_ast::Expr;
-
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::Expr;
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -47,6 +47,10 @@ impl Violation for SixPY3 {
 
 /// YTT202
 pub(crate) fn name_or_attribute(checker: &mut Checker, expr: &Expr) {
+    if !checker.semantic().seen_module(Modules::SIX) {
+        return;
+    }
+
     if checker
         .semantic()
         .resolve_call_path(expr)

diff --git a/crates/ruff_linter/src/rules/flake8_async/rules/blocking_os_call.rs b/crates/ruff_linter/src/rules/flake8_async/rules/blocking_os_call.rs
@@ -3,6 +3,7 @@ use ruff_python_ast::ExprCall;
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::call_path::CallPath;
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -42,17 +43,19 @@ impl Violation for BlockingOsCallInAsyncFunction {
 
 /// ASYNC102
 pub(crate) fn blocking_os_call(checker: &mut Checker, call: &ExprCall) {
-    if checker.semantic().in_async_context() {
-        if checker
-            .semantic()
-            .resolve_call_path(call.func.as_ref())
-            .as_ref()
-            .is_some_and(is_unsafe_os_method)
-        {
-            checker.diagnostics.push(Diagnostic::new(
-                BlockingOsCallInAsyncFunction,
-                call.func.range(),
-            ));
+    if checker.semantic().seen_module(Modules::OS) {
+        if checker.semantic().in_async_context() {
+            if checker
+                .semantic()
+                .resolve_call_path(call.func.as_ref())
+                .as_ref()
+                .is_some_and(is_unsafe_os_method)
+            {
+                checker.diagnostics.push(Diagnostic::new(
+                    BlockingOsCallInAsyncFunction,
+                    call.func.range(),
+                ));
+            }
         }
     }
 }

diff --git a/crates/ruff_linter/src/rules/flake8_bandit/rules/bad_file_permissions.rs b/crates/ruff_linter/src/rules/flake8_bandit/rules/bad_file_permissions.rs
@@ -4,7 +4,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::call_path::CallPath;
 use ruff_python_ast::{self as ast, Expr, Operator};
-use ruff_python_semantic::SemanticModel;
+use ruff_python_semantic::{Modules, SemanticModel};
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -60,6 +60,10 @@ enum Reason {
 
 /// S103
 pub(crate) fn bad_file_permissions(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::OS) {
+        return;
+    }
+
     if checker
         .semantic()
         .resolve_call_path(&call.func)

diff --git a/crates/ruff_linter/src/rules/flake8_bandit/rules/django_raw_sql.rs b/crates/ruff_linter/src/rules/flake8_bandit/rules/django_raw_sql.rs
@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast, Expr};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -35,6 +36,10 @@ impl Violation for DjangoRawSql {
 
 /// S611
 pub(crate) fn django_raw_sql(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::DJANGO) {
+        return;
+    }
+
     if checker
         .semantic()
         .resolve_call_path(&call.func)

diff --git a/crates/ruff_linter/src/rules/flake8_bandit/rules/logging_config_insecure_listen.rs b/crates/ruff_linter/src/rules/flake8_bandit/rules/logging_config_insecure_listen.rs
@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -35,6 +36,10 @@ impl Violation for LoggingConfigInsecureListen {
 
 /// S612
 pub(crate) fn logging_config_insecure_listen(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::LOGGING) {
+        return;
+    }
+
     if checker
         .semantic()
         .resolve_call_path(&call.func)

diff --git a/crates/ruff_linter/src/rules/flake8_bandit/rules/tarfile_unsafe_members.rs b/crates/ruff_linter/src/rules/flake8_bandit/rules/tarfile_unsafe_members.rs
@@ -3,6 +3,7 @@ use ruff_diagnostics::Diagnostic;
 use ruff_diagnostics::Violation;
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 /// ## What it does
@@ -48,6 +49,10 @@ impl Violation for TarfileUnsafeMembers {
 
 /// S202
 pub(crate) fn tarfile_unsafe_members(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::TARFILE) {
+        return;
+    }
+
     if !call
         .func
         .as_attribute_expr()
@@ -65,10 +70,6 @@ pub(crate) fn tarfile_unsafe_members(checker: &mut Checker, call: &ast::ExprCall
         return;
     }
 
-    if !checker.semantic().seen(&["tarfile"]) {
-        return;
-    }
-
     checker
         .diagnostics
         .push(Diagnostic::new(TarfileUnsafeMembers, call.func.range()));

diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/rules/re_sub_positional_args.rs b/crates/ruff_linter/src/rules/flake8_bugbear/rules/re_sub_positional_args.rs
@@ -4,6 +4,7 @@ use ruff_python_ast::{self as ast};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -56,6 +57,10 @@ impl Violation for ReSubPositionalArgs {
 
 /// B034
 pub(crate) fn re_sub_positional_args(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::RE) {
+        return;
+    }
+
     let Some(method) = checker
         .semantic()
         .resolve_call_path(&call.func)

diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_date_fromtimestamp.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_date_fromtimestamp.rs
@@ -3,6 +3,7 @@ use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::Modules;
 
 use crate::checkers::ast::Checker;
 
@@ -57,6 +58,10 @@ impl Violation for CallDateFromtimestamp {
 }
 
 pub(crate) fn call_date_fromtimestamp(checker: &mut Checker, func: &Expr, location: TextRange) {
+    if !checker.semantic().seen_module(Modules::DATETIME) {
+        return;
+    }
+
     if checker
         .semantic()
         .resolve_call_path(func)

diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_date_today.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_date_today.rs
@@ -3,6 +3,7 @@ use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::Modules;
 
 use crate::checkers::ast::Checker;
 
@@ -56,6 +57,10 @@ impl Violation for CallDateToday {
 }
 
 pub(crate) fn call_date_today(checker: &mut Checker, func: &Expr, location: TextRange) {
+    if !checker.semantic().seen_module(Modules::DATETIME) {
+        return;
+    }
+
     if checker
         .semantic()
         .resolve_call_path(func)

diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_fromtimestamp.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_fromtimestamp.rs
@@ -2,6 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 
 use ruff_python_ast::{self as ast};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -60,6 +61,10 @@ impl Violation for CallDatetimeFromtimestamp {
 }
 
 pub(crate) fn call_datetime_fromtimestamp(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::DATETIME) {
+        return;
+    }
+
     if !checker
         .semantic()
         .resolve_call_path(&call.func)

diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_now_without_tzinfo.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_now_without_tzinfo.rs
@@ -2,6 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 
 use ruff_python_ast::{self as ast, Expr};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -56,6 +57,10 @@ impl Violation for CallDatetimeNowWithoutTzinfo {
 }
 
 pub(crate) fn call_datetime_now_without_tzinfo(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::DATETIME) {
+        return;
+    }
+
     if !checker
         .semantic()
         .resolve_call_path(&call.func)

diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_strptime_without_zone.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_strptime_without_zone.rs
@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast, Expr};
+use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -64,6 +65,10 @@ impl Violation for CallDatetimeStrptimeWithoutZone {
 
 /// DTZ007
 pub(crate) fn call_datetime_strptime_without_zone(checker: &mut Checker, call: &ast::ExprCall) {
+    if !checker.semantic().seen_module(Modules::DATETIME) {
+        return;
+    }
+
     if !checker
         .semantic()
         .resolve_call_path(&call.func)

diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_today.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_today.rs
@@ -3,6 +3,7 @@ use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::Modules;
 
 use crate::checkers::ast::Checker;
 
@@ -55,6 +56,10 @@ impl Violation for CallDatetimeToday {
 }
 
 pub(crate) fn call_datetime_today(checker: &mut Checker, func: &Expr, location: TextRange) {
+    if !checker.semantic().seen_module(Modules::DATETIME) {
+        return;
+    }
+
     if !checker
         .semantic()
         .resolve_call_path(func)

diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_utcfromtimestamp.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_utcfromtimestamp.rs
@@ -3,6 +3,7 @@ use ruff_text_size::TextRange;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::Modules;
 
 use crate::checkers::ast::Checker;
 
@@ -63,6 +64,10 @@ pub(crate) fn call_datetime_utcfromtimestamp(
     func: &Expr,
     location: TextRange,
 ) {
+    if !checker.semantic().seen_module(Modules::DATETIME) {
+        return;
+    }
+
     if !checker
         .semantic()
         .resolve_call_path(func)