From b49c80f8c8b856997bfe2851ce6a298da2ca8f4c Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sat, 12 Aug 2023 14:52:44 -0400
Subject: [PATCH] Use top-level semantic detection for E402 (#6526)

## Summary

Noticed in https://github.com/astral-sh/ruff/pull/6378. Given `import h;
import i`, we don't consider `import i` to be a "top-level" import for
E402 purposes, which is wrong. Similarly, we _do_ consider `import k` to
be a "top-level" import in:

```python
if __name__ == "__main__":
    import j; \
import k
```

Using the semantic detection, rather than relying on newline position,
fixes both cases.

## Test Plan

`cargo test`
---
 .../resources/test/fixtures/pycodestyle/E402.py  |  7 +++++++
 .../ruff/src/checkers/ast/analyze/statement.rs   | 12 ++----------
 .../ruff/src/rules/pycodestyle/rules/imports.rs  | 12 +++---------
 ..._rules__pycodestyle__tests__E402_E402.py.snap | 16 ++++++++++++++++
 4 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/crates/ruff/resources/test/fixtures/pycodestyle/E402.py b/crates/ruff/resources/test/fixtures/pycodestyle/E402.py
index 81e6306c7af11..d5943407852a2 100644
--- a/crates/ruff/resources/test/fixtures/pycodestyle/E402.py
+++ b/crates/ruff/resources/test/fixtures/pycodestyle/E402.py
@@ -30,3 +30,10 @@ def foo() -> None:
 
 if __name__ == "__main__":
     import g
+
+import h; import i
+
+
+if __name__ == "__main__":
+    import j; \
+import k
diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs
index 81a36c87d6c76..ada91344da020 100644
--- a/crates/ruff/src/checkers/ast/analyze/statement.rs
+++ b/crates/ruff/src/checkers/ast/analyze/statement.rs
@@ -520,11 +520,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 pycodestyle::rules::multiple_imports_on_one_line(checker, stmt, names);
             }
             if checker.enabled(Rule::ModuleImportNotAtTopOfFile) {
-                pycodestyle::rules::module_import_not_at_top_of_file(
-                    checker,
-                    stmt,
-                    checker.locator,
-                );
+                pycodestyle::rules::module_import_not_at_top_of_file(checker, stmt);
             }
             if checker.enabled(Rule::GlobalStatement) {
                 for name in names {
@@ -689,11 +685,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             let module = module.as_deref();
             let level = level.map(|level| level.to_u32());
             if checker.enabled(Rule::ModuleImportNotAtTopOfFile) {
-                pycodestyle::rules::module_import_not_at_top_of_file(
-                    checker,
-                    stmt,
-                    checker.locator,
-                );
+                pycodestyle::rules::module_import_not_at_top_of_file(checker, stmt);
             }
             if checker.enabled(Rule::GlobalStatement) {
                 for name in names {
diff --git a/crates/ruff/src/rules/pycodestyle/rules/imports.rs b/crates/ruff/src/rules/pycodestyle/rules/imports.rs
index 4085b6aa66672..8e20077d34ac3 100644
--- a/crates/ruff/src/rules/pycodestyle/rules/imports.rs
+++ b/crates/ruff/src/rules/pycodestyle/rules/imports.rs
@@ -1,8 +1,6 @@
-use ruff_python_ast::{Alias, Ranged, Stmt};
-
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_source_file::Locator;
+use ruff_python_ast::{Alias, Ranged, Stmt};
 
 use crate::checkers::ast::Checker;
 
@@ -81,12 +79,8 @@ pub(crate) fn multiple_imports_on_one_line(checker: &mut Checker, stmt: &Stmt, n
 }
 
 /// E402
-pub(crate) fn module_import_not_at_top_of_file(
-    checker: &mut Checker,
-    stmt: &Stmt,
-    locator: &Locator,
-) {
-    if checker.semantic().seen_import_boundary() && locator.is_at_start_of_line(stmt.start()) {
+pub(crate) fn module_import_not_at_top_of_file(checker: &mut Checker, stmt: &Stmt) {
+    if checker.semantic().seen_import_boundary() && checker.semantic().at_top_level() {
         checker
             .diagnostics
             .push(Diagnostic::new(ModuleImportNotAtTopOfFile, stmt.range()));
diff --git a/crates/ruff/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E402_E402.py.snap b/crates/ruff/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E402_E402.py.snap
index f768015c97841..a0518e1437172 100644
--- a/crates/ruff/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E402_E402.py.snap
+++ b/crates/ruff/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E402_E402.py.snap
@@ -9,4 +9,20 @@ E402.py:24:1: E402 Module level import not at top of file
    | ^^^^^^^^ E402
    |
 
+E402.py:34:1: E402 Module level import not at top of file
+   |
+32 |     import g
+33 | 
+34 | import h; import i
+   | ^^^^^^^^ E402
+   |
+
+E402.py:34:11: E402 Module level import not at top of file
+   |
+32 |     import g
+33 | 
+34 | import h; import i
+   |           ^^^^^^^^ E402
+   |
+