From c28ac755910dcd0afcae4e143ae454aded270aa2 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Wed, 31 Aug 2022 19:11:26 -0400
Subject: [PATCH] Implement F841 (for functions) (#65)

---
 resources/test/src/F841.py |  6 ++++++
 src/check_ast.rs           | 31 ++++++++++++++++++++++++++++---
 src/linter.rs              | 17 ++++++++++++-----
 3 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/resources/test/src/F841.py b/resources/test/src/F841.py
index d454a80df6..c343debf1e 100644
--- a/resources/test/src/F841.py
+++ b/resources/test/src/F841.py
@@ -8,3 +8,9 @@ try:
     1 / 0
 except ValueError as e:
     print(e)
+
+
+def f():
+    x = 1
+    y = 2
+    z = x + y
diff --git a/src/check_ast.rs b/src/check_ast.rs
index f7a9fc38f6..52054c4eca 100644
--- a/src/check_ast.rs
+++ b/src/check_ast.rs
@@ -287,9 +287,34 @@ impl Visitor for Checker<'_> {
         visitor::walk_stmt(self, stmt);
 
         match &stmt.node {
-            StmtKind::ClassDef { .. }
-            | StmtKind::FunctionDef { .. }
-            | StmtKind::AsyncFunctionDef { .. } => self.pop_scope(),
+            StmtKind::ClassDef { .. } => {
+                if let Some(scope) = self.scopes.pop() {
+                    self.dead_scopes.push(scope);
+                }
+            }
+            StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+                let scope = self.scopes.last().expect("No current scope found.");
+                for (name, binding) in scope.values.iter() {
+                    // TODO(charlie): Ignore if using `locals`.
+                    if self.settings.select.contains(&CheckCode::F841)
+                        && binding.used.is_none()
+                        && name != "_"
+                        && name != "__tracebackhide__"
+                        && name != "__traceback_info__"
+                        && name != "__traceback_supplement__"
+                        && matches!(binding.kind, BindingKind::Assignment)
+                    {
+                        self.checks.push(Check {
+                            kind: CheckKind::UnusedVariable(name.to_string()),
+                            location: binding.location,
+                        });
+                    }
+                }
+
+                if let Some(scope) = self.scopes.pop() {
+                    self.dead_scopes.push(scope);
+                }
+            }
             _ => {}
         };
 
diff --git a/src/linter.rs b/src/linter.rs
index d59e1feccc..76633073b3 100644
--- a/src/linter.rs
+++ b/src/linter.rs
@@ -365,11 +365,18 @@ mod tests {
             },
             &cache::Mode::None,
         )?;
-        let expected = vec![Message {
-            kind: CheckKind::UnusedVariable("e".to_string()),
-            location: Location::new(3, 1),
-            filename: "./resources/test/src/F841.py".to_string(),
-        }];
+        let expected = vec![
+            Message {
+                kind: CheckKind::UnusedVariable("e".to_string()),
+                location: Location::new(3, 1),
+                filename: "./resources/test/src/F841.py".to_string(),
+            },
+            Message {
+                kind: CheckKind::UnusedVariable("z".to_string()),
+                location: Location::new(16, 5),
+                filename: "./resources/test/src/F841.py".to_string(),
+            },
+        ];
         assert_eq!(actual.len(), expected.len());
         for i in 0..actual.len() {
             assert_eq!(actual[i], expected[i]);