Add scope tracking (and F706) (#22)

Cargo.lock

@@ -1534,7 +1534,7 @@ dependencies = [
 
 [[package]]
 name = "rust-python-linter"
-version = "0.0.10"
+version = "0.0.11"
 dependencies = [
  "anyhow",
  "bincode",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rust-python-linter"
-version = "0.0.10"
+version = "0.0.11"
 edition = "2021"
 
 [lib]

resources/test/src/return_outside_function.py

@@ -0,0 +1,9 @@
+def f() -> int:
+    return 1
+
+
+class Foo:
+    return 2
+
+
+return 3

src/check_ast.rs

@@ -1,5 +1,6 @@
 use std::collections::BTreeSet;
 
+use crate::check_ast::ScopeKind::{Class, Function, Generator, Module};
 use rustpython_parser::ast::{Arg, Arguments, Expr, ExprKind, Stmt, StmtKind, Suite};
 
 use crate::checks::{Check, CheckKind};
@@ -7,9 +8,21 @@ use crate::settings::Settings;
 use crate::visitor;
 use crate::visitor::Visitor;
 
+enum ScopeKind {
+    Class,
+    Function,
+    Generator,
+    Module,
+}
+
+struct Scope {
+    kind: ScopeKind,
+}
+
 struct Checker<'a> {
     settings: &'a Settings,
     checks: Vec<Check>,
+    scopes: Vec<Scope>,
 }
 
 impl Checker<'_> {
@@ -17,6 +30,7 @@ impl Checker<'_> {
         Checker {
             settings,
             checks: vec![],
+            scopes: vec![Scope { kind: Module }],
         }
     }
 }
@@ -24,6 +38,28 @@ impl Checker<'_> {
 impl Visitor for Checker<'_> {
     fn visit_stmt(&mut self, stmt: &Stmt) {
         match &stmt.node {
+            StmtKind::FunctionDef { .. } => self.scopes.push(Scope { kind: Function }),
+            StmtKind::AsyncFunctionDef { .. } => self.scopes.push(Scope { kind: Function }),
+            StmtKind::Return { .. } => {
+                if self
+                    .settings
+                    .select
+                    .contains(CheckKind::ReturnOutsideFunction.code())
+                {
+                    if let Some(scope) = self.scopes.last() {
+                        match scope.kind {
+                            Class | Module => {
+                                self.checks.push(Check {
+                                    kind: CheckKind::ReturnOutsideFunction,
+                                    location: stmt.location,
+                                });
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+            }
+            StmtKind::ClassDef { .. } => self.scopes.push(Scope { kind: Class }),
             StmtKind::ImportFrom { names, .. } => {
                 if self
                     .settings
@@ -85,18 +121,29 @@ impl Visitor for Checker<'_> {
         }
 
         visitor::walk_stmt(self, stmt);
+
+        match &stmt.node {
+            StmtKind::ClassDef { .. }
+            | StmtKind::FunctionDef { .. }
+            | StmtKind::AsyncFunctionDef { .. } => {
+                self.scopes.pop();
+            }
+            _ => {}
+        };
     }
 
     fn visit_expr(&mut self, expr: &Expr) {
-        if self
-            .settings
-            .select
-            .contains(CheckKind::FStringMissingPlaceholders.code())
-        {
-            if let ExprKind::JoinedStr { values } = &expr.node {
-                if !values
-                    .iter()
-                    .any(|value| matches!(value.node, ExprKind::FormattedValue { .. }))
+        match &expr.node {
+            ExprKind::GeneratorExp { .. } => self.scopes.push(Scope { kind: Generator }),
+            ExprKind::Lambda { .. } => self.scopes.push(Scope { kind: Function }),
+            ExprKind::JoinedStr { values } => {
+                if self
+                    .settings
+                    .select
+                    .contains(CheckKind::FStringMissingPlaceholders.code())
+                    && !values
+                        .iter()
+                        .any(|value| matches!(value.node, ExprKind::FormattedValue { .. }))
                 {
                     self.checks.push(Check {
                         kind: CheckKind::FStringMissingPlaceholders,
@@ -104,9 +151,17 @@ impl Visitor for Checker<'_> {
                     });
                 }
             }
-        }
+            _ => {}
+        };
 
         visitor::walk_expr(self, expr);
+
+        match &expr.node {
+            ExprKind::GeneratorExp { .. } | ExprKind::Lambda { .. } => {
+                self.scopes.pop();
+            }
+            _ => {}
+        };
     }
 
     fn visit_arguments(&mut self, arguments: &Arguments) {

src/checks.rs

@@ -7,6 +7,7 @@ pub enum CheckCode {
     F541,
     F634,
     F403,
+    F706,
     F901,
     E501,
 }
@@ -18,6 +19,7 @@ impl CheckCode {
             CheckCode::F541 => "F541",
             CheckCode::F634 => "F634",
             CheckCode::F403 => "F403",
+            CheckCode::F706 => "F706",
             CheckCode::F901 => "F901",
             CheckCode::E501 => "E501",
         }
@@ -36,6 +38,7 @@ pub enum CheckKind {
     FStringMissingPlaceholders,
     IfTuple,
     ImportStarUsage,
+    ReturnOutsideFunction,
     RaiseNotImplemented,
     LineTooLong,
 }
@@ -48,6 +51,7 @@ impl CheckKind {
             CheckCode::F541 => CheckKind::FStringMissingPlaceholders,
             CheckCode::F634 => CheckKind::IfTuple,
             CheckCode::F403 => CheckKind::ImportStarUsage,
+            CheckCode::F706 => CheckKind::ReturnOutsideFunction,
             CheckCode::F901 => CheckKind::RaiseNotImplemented,
             CheckCode::E501 => CheckKind::LineTooLong,
         }
@@ -60,6 +64,7 @@ impl CheckKind {
             CheckKind::FStringMissingPlaceholders => &CheckCode::F541,
             CheckKind::IfTuple => &CheckCode::F634,
             CheckKind::ImportStarUsage => &CheckCode::F403,
+            CheckKind::ReturnOutsideFunction => &CheckCode::F706,
             CheckKind::RaiseNotImplemented => &CheckCode::F901,
             CheckKind::LineTooLong => &CheckCode::E501,
         }
@@ -72,6 +77,7 @@ impl CheckKind {
             CheckKind::FStringMissingPlaceholders => "f-string without any placeholders",
             CheckKind::IfTuple => "If test is a tuple, which is always `True`",
             CheckKind::ImportStarUsage => "Unable to detect undefined names",
+            CheckKind::ReturnOutsideFunction => "a `return` statement outside of a function/method",
             CheckKind::RaiseNotImplemented => {
                 "'raise NotImplemented' should be 'raise NotImplementedError"
             }
@@ -86,6 +92,7 @@ impl CheckKind {
             CheckKind::FStringMissingPlaceholders => &LintSource::AST,
             CheckKind::IfTuple => &LintSource::AST,
             CheckKind::ImportStarUsage => &LintSource::AST,
+            CheckKind::ReturnOutsideFunction => &LintSource::AST,
             CheckKind::RaiseNotImplemented => &LintSource::AST,
             CheckKind::LineTooLong => &LintSource::Lines,
         }

src/linter.rs

@@ -205,6 +205,37 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn return_outside_function() -> Result<()> {
+        let actual = check_path(
+            &Path::new("./resources/test/src/return_outside_function.py"),
+            &settings::Settings {
+                line_length: 88,
+                exclude: vec![],
+                select: BTreeSet::from([CheckCode::F706]),
+            },
+            &cache::Mode::None,
+        )?;
+        let expected = vec![
+            Message {
+                kind: CheckKind::ReturnOutsideFunction,
+                location: Location::new(6, 5),
+                filename: "./resources/test/src/return_outside_function.py".to_string(),
+            },
+            Message {
+                kind: CheckKind::ReturnOutsideFunction,
+                location: Location::new(9, 1),
+                filename: "./resources/test/src/return_outside_function.py".to_string(),
+            },
+        ];
+        assert_eq!(actual.len(), expected.len());
+        for i in 0..actual.len() {
+            assert_eq!(actual[i], expected[i]);
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn raise_not_implemented() -> Result<()> {
         let actual = check_path(

src/settings.rs

@@ -45,6 +45,7 @@ impl Settings {
                     CheckCode::F541,
                     CheckCode::F634,
                     CheckCode::F403,
+                    CheckCode::F706,
                     CheckCode::F901,
                     CheckCode::E501,
                 ])