Implement E402 (ModuleImportNotAtTopOfFile) (#102)

README.md

@@ -110,6 +110,7 @@ OPTIONS:
 
 | Code | Name | Message |
 | ---- | ----- | ------- |
+| E402 | ModuleImportNotAtTopOfFile | Module level import not at top of file |
 | E501 | LineTooLong | Line too long |
 | F401 | UnusedImport | `...` imported but unused |
 | F403 | ImportStarUsage | Unable to detect undefined names |

examples/generate_rules_table.rs

@@ -10,6 +10,7 @@ fn main() {
         CheckKind::IfTuple,
         CheckKind::ImportStarUsage,
         CheckKind::LineTooLong,
+        CheckKind::ModuleImportNotAtTopOfFile,
         CheckKind::RaiseNotImplemented,
         CheckKind::ReturnOutsideFunction,
         CheckKind::UndefinedExport("...".to_string()),

resources/test/fixtures/E402.py

@@ -0,0 +1,28 @@
+"""Top-level docstring."""
+import a
+
+try:
+    import b
+except ImportError:
+    pass
+else:
+    pass
+
+import c
+
+if x > 0:
+    import d
+else:
+    import e
+
+y = x + 1
+
+import f
+
+
+def foo() -> None:
+    import e
+
+
+if __name__ == "__main__":
+    import g

resources/test/fixtures/pyproject.toml

@@ -2,6 +2,7 @@
 line-length = 88
 exclude = ["excluded.py", "**/migrations"]
 select = [
+    "E402",
     "E501",
     "F401",
     "F403",

src/check_ast.rs

@@ -23,6 +23,8 @@ struct Checker<'a> {
     deferred: Vec<String>,
     in_f_string: bool,
     in_annotation: bool,
+    seen_non_import: bool,
+    seen_docstring: bool,
 }
 
 impl Checker<'_> {
@@ -36,6 +38,8 @@ impl Checker<'_> {
             deferred: vec![],
             in_f_string: false,
             in_annotation: false,
+            seen_non_import: false,
+            seen_docstring: false,
         }
     }
 }
@@ -154,6 +158,19 @@ impl Visitor for Checker<'_> {
                 self.push_scope(Scope::new(ScopeKind::Class))
             }
             StmtKind::Import { names } => {
+                if self
+                    .settings
+                    .select
+                    .contains(CheckKind::ModuleImportNotAtTopOfFile.code())
+                    && self.seen_non_import
+                    && stmt.location.column() == 1
+                {
+                    self.checks.push(Check {
+                        kind: CheckKind::ModuleImportNotAtTopOfFile,
+                        location: stmt.location,
+                    });
+                }
+
                 for alias in names {
                     if alias.node.name.contains('.') && alias.node.asname.is_none() {
                         // TODO(charlie): Multiple submodule imports with the same parent module
@@ -191,6 +208,19 @@ impl Visitor for Checker<'_> {
                 }
             }
             StmtKind::ImportFrom { names, module, .. } => {
+                if self
+                    .settings
+                    .select
+                    .contains(CheckKind::ModuleImportNotAtTopOfFile.code())
+                    && self.seen_non_import
+                    && stmt.location.column() == 1
+                {
+                    self.checks.push(Check {
+                        kind: CheckKind::ModuleImportNotAtTopOfFile,
+                        location: stmt.location,
+                    });
+                }
+
                 for alias in names {
                     let name = alias
                         .node
@@ -282,8 +312,12 @@ impl Visitor for Checker<'_> {
                     }
                 }
             }
-            StmtKind::AugAssign { target, .. } => self.handle_node_load(target),
+            StmtKind::AugAssign { target, .. } => {
+                self.seen_non_import = true;
+                self.handle_node_load(target);
+            }
             StmtKind::Assert { test, .. } => {
+                self.seen_non_import = true;
                 if self.settings.select.contains(CheckKind::AssertTuple.code()) {
                     if let ExprKind::Tuple { elts, .. } = &test.node {
                         if !elts.is_empty() {
@@ -312,6 +346,22 @@ impl Visitor for Checker<'_> {
                     }
                 }
             }
+            StmtKind::Expr { value } => {
+                if !self.seen_docstring {
+                    if let ExprKind::Constant {
+                        value: Constant::Str(_),
+                        ..
+                    } = &value.node
+                    {
+                        self.seen_docstring = true;
+                    }
+                } else {
+                    self.seen_non_import = true;
+                }
+            }
+            StmtKind::Delete { .. } | StmtKind::Assign { .. } | StmtKind::AnnAssign { .. } => {
+                self.seen_non_import = true;
+            }
             _ => {}
         }
 

src/checks.rs

@@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};
 
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Hash, PartialOrd, Ord)]
 pub enum CheckCode {
+    E402,
     E501,
     F401,
     F403,
@@ -31,6 +32,7 @@ impl FromStr for CheckCode {
 
     fn from_str(s: &str) -> Result<Self> {
         match s {
+            "E402" => Ok(CheckCode::E402),
             "E501" => Ok(CheckCode::E501),
             "F401" => Ok(CheckCode::F401),
             "F403" => Ok(CheckCode::F403),
@@ -55,6 +57,7 @@ impl FromStr for CheckCode {
 impl CheckCode {
     pub fn as_str(&self) -> &str {
         match self {
+            CheckCode::E402 => "E402",
             CheckCode::E501 => "E501",
             CheckCode::F401 => "F401",
             CheckCode::F403 => "F403",
@@ -77,6 +80,7 @@ impl CheckCode {
     /// The source for the check (either the AST, or the physical lines).
     pub fn lint_source(&self) -> &'static LintSource {
         match self {
+            CheckCode::E402 => &LintSource::AST,
             CheckCode::E501 => &LintSource::Lines,
             CheckCode::F401 => &LintSource::AST,
             CheckCode::F403 => &LintSource::AST,
@@ -105,17 +109,18 @@ pub enum LintSource {
 
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub enum CheckKind {
+    AssertTuple,
+    DefaultExceptNotLast,
     DuplicateArgumentName,
     FStringMissingPlaceholders,
-    AssertTuple,
     IfTuple,
     ImportStarUsage,
     LineTooLong,
+    ModuleImportNotAtTopOfFile,
     RaiseNotImplemented,
     ReturnOutsideFunction,
-    DefaultExceptNotLast,
-    UndefinedLocal(String),
     UndefinedExport(String),
+    UndefinedLocal(String),
     UndefinedName(String),
     UnusedImport(String),
     UnusedVariable(String),
@@ -127,17 +132,18 @@ impl CheckKind {
     /// The name of the check.
     pub fn name(&self) -> &'static str {
         match self {
+            CheckKind::AssertTuple => "AssertTuple",
+            CheckKind::DefaultExceptNotLast => "DefaultExceptNotLast",
             CheckKind::DuplicateArgumentName => "DuplicateArgumentName",
             CheckKind::FStringMissingPlaceholders => "FStringMissingPlaceholders",
-            CheckKind::AssertTuple => "AssertTuple",
             CheckKind::IfTuple => "IfTuple",
             CheckKind::ImportStarUsage => "ImportStarUsage",
             CheckKind::LineTooLong => "LineTooLong",
+            CheckKind::ModuleImportNotAtTopOfFile => "ModuleImportNotAtTopOfFile",
             CheckKind::RaiseNotImplemented => "RaiseNotImplemented",
             CheckKind::ReturnOutsideFunction => "ReturnOutsideFunction",
-            CheckKind::DefaultExceptNotLast => "DefaultExceptNotLast",
-            CheckKind::UndefinedLocal(_) => "UndefinedLocal",
             CheckKind::UndefinedExport(_) => "UndefinedExport",
+            CheckKind::UndefinedLocal(_) => "UndefinedLocal",
             CheckKind::UndefinedName(_) => "UndefinedName",
             CheckKind::UnusedImport(_) => "UnusedImport",
             CheckKind::UnusedVariable(_) => "UnusedVariable",
@@ -149,15 +155,16 @@ impl CheckKind {
     /// A four-letter shorthand code for the check.
     pub fn code(&self) -> &'static CheckCode {
         match self {
+            CheckKind::AssertTuple => &CheckCode::F631,
+            CheckKind::DefaultExceptNotLast => &CheckCode::F707,
             CheckKind::DuplicateArgumentName => &CheckCode::F831,
             CheckKind::FStringMissingPlaceholders => &CheckCode::F541,
-            CheckKind::AssertTuple => &CheckCode::F631,
             CheckKind::IfTuple => &CheckCode::F634,
             CheckKind::ImportStarUsage => &CheckCode::F403,
             CheckKind::LineTooLong => &CheckCode::E501,
+            CheckKind::ModuleImportNotAtTopOfFile => &CheckCode::E402,
             CheckKind::RaiseNotImplemented => &CheckCode::F901,
             CheckKind::ReturnOutsideFunction => &CheckCode::F706,
-            CheckKind::DefaultExceptNotLast => &CheckCode::F707,
             CheckKind::UndefinedExport(_) => &CheckCode::F822,
             CheckKind::UndefinedLocal(_) => &CheckCode::F823,
             CheckKind::UndefinedName(_) => &CheckCode::F821,
@@ -171,27 +178,30 @@ impl CheckKind {
     /// The body text for the check.
     pub fn body(&self) -> String {
         match self {
+            CheckKind::AssertTuple => {
+                "Assert test is a non-empty tuple, which is always `True`".to_string()
+            }
+            CheckKind::DefaultExceptNotLast => {
+                "an `except:` block as not the last exception handler".to_string()
+            }
             CheckKind::DuplicateArgumentName => {
                 "Duplicate argument name in function definition".to_string()
             }
             CheckKind::FStringMissingPlaceholders => {
                 "f-string without any placeholders".to_string()
             }
-            CheckKind::AssertTuple => {
-                "Assert test is a non-empty tuple, which is always `True`".to_string()
-            }
             CheckKind::IfTuple => "If test is a tuple, which is always `True`".to_string(),
             CheckKind::ImportStarUsage => "Unable to detect undefined names".to_string(),
             CheckKind::LineTooLong => "Line too long".to_string(),
+            CheckKind::ModuleImportNotAtTopOfFile => {
+                "Module level import not at top of file".to_string()
+            }
             CheckKind::RaiseNotImplemented => {
                 "`raise NotImplemented` should be `raise NotImplementedError`".to_string()
             }
             CheckKind::ReturnOutsideFunction => {
                 "a `return` statement outside of a function/method".to_string()
             }
-            CheckKind::DefaultExceptNotLast => {
-                "an `except:` block as not the last exception handler".to_string()
-            }
             CheckKind::UndefinedExport(name) => {
                 format!("Undefined name `{name}` in `__all__`")
             }

src/linter.rs

@@ -65,6 +65,30 @@ mod tests {
     use crate::message::Message;
     use crate::{cache, settings};
 
+    #[test]
+    fn e402() -> Result<()> {
+        let actual = check_path(
+            Path::new("./resources/test/fixtures/E402.py"),
+            &settings::Settings {
+                line_length: 88,
+                exclude: vec![],
+                select: BTreeSet::from([CheckCode::E402]),
+            },
+            &cache::Mode::None,
+        )?;
+        let expected = vec![Message {
+            kind: CheckKind::ModuleImportNotAtTopOfFile,
+            location: Location::new(20, 1),
+            filename: "./resources/test/fixtures/E402.py".to_string(),
+        }];
+        assert_eq!(actual.len(), expected.len());
+        for i in 0..actual.len() {
+            assert_eq!(actual[i], expected[i]);
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn e501() -> Result<()> {
         let actual = check_path(

src/pyproject.rs

@@ -237,6 +237,7 @@ other-attribute = 1
                     Path::new("**/migrations").to_path_buf()
                 ]),
                 select: Some(BTreeSet::from([
+                    CheckCode::E402,
                     CheckCode::E501,
                     CheckCode::F401,
                     CheckCode::F403,

src/settings.rs

@@ -44,6 +44,7 @@ impl Settings {
                 .collect(),
             select: config.select.unwrap_or_else(|| {
                 BTreeSet::from([
+                    CheckCode::E402,
                     CheckCode::E501,
                     CheckCode::F401,
                     CheckCode::F403,
@@ -52,8 +53,8 @@ impl Settings {
                     CheckCode::F634,
                     CheckCode::F706,
                     CheckCode::F707,
-                    CheckCode::F831,
                     CheckCode::F823,
+                    CheckCode::F831,
                     CheckCode::F901,
                 ])
             }),