Upgrade parser to handle more F821 cases (#57)

Cargo.lock

@@ -1662,7 +1662,7 @@ dependencies = [
 [[package]]
 name = "rustpython-ast"
 version = "0.1.0"
-source = "git+https://github.com/charliermarsh/RustPython.git?rev=97a9d9de687227595932d6c8a04014dcfc892ff4#97a9d9de687227595932d6c8a04014dcfc892ff4"
+source = "git+https://github.com/charliermarsh/RustPython.git?rev=1613f6c6990011a4bc559e79aaf28d715f9f729b#1613f6c6990011a4bc559e79aaf28d715f9f729b"
 dependencies = [
  "num-bigint",
  "rustpython-compiler-core",
@@ -1671,7 +1671,7 @@ dependencies = [
 [[package]]
 name = "rustpython-compiler-core"
 version = "0.1.2"
-source = "git+https://github.com/charliermarsh/RustPython.git?rev=97a9d9de687227595932d6c8a04014dcfc892ff4#97a9d9de687227595932d6c8a04014dcfc892ff4"
+source = "git+https://github.com/charliermarsh/RustPython.git?rev=1613f6c6990011a4bc559e79aaf28d715f9f729b#1613f6c6990011a4bc559e79aaf28d715f9f729b"
 dependencies = [
  "bincode",
  "bitflags",
@@ -1688,7 +1688,7 @@ dependencies = [
 [[package]]
 name = "rustpython-parser"
 version = "0.1.2"
-source = "git+https://github.com/charliermarsh/RustPython.git?rev=97a9d9de687227595932d6c8a04014dcfc892ff4#97a9d9de687227595932d6c8a04014dcfc892ff4"
+source = "git+https://github.com/charliermarsh/RustPython.git?rev=1613f6c6990011a4bc559e79aaf28d715f9f729b#1613f6c6990011a4bc559e79aaf28d715f9f729b"
 dependencies = [
  "ahash",
  "anyhow",

Cargo.toml

@@ -22,7 +22,7 @@ log = { version = "0.4.17" }
 notify = { version = "4.0.17" }
 rayon = { version = "1.5.3" }
 regex = { version = "1.6.0" }
-rustpython-parser = { features = ["lalrpop"], git = "https://github.com/charliermarsh/RustPython.git", rev = "97a9d9de687227595932d6c8a04014dcfc892ff4" }
+rustpython-parser = { features = ["lalrpop"], git = "https://github.com/charliermarsh/RustPython.git", rev = "1613f6c6990011a4bc559e79aaf28d715f9f729b" }
 serde = { version = "1.0.143", features = ["derive"] }
 serde_json = { version = "1.0.83" }
 toml = { version = "0.5.9" }

resources/test/src/F821.py

@@ -13,4 +13,36 @@ def get_name():
 def get_name(self):
     return self.name
 
+
 x = list()
+
+
+def randdec(maxprec, maxexp):
+    return numeric_string(maxprec, maxexp)
+
+
+def ternary_optarg(prec, exp_range, itr):
+    for _ in range(100):
+        a = randdec(prec, 2 * exp_range)
+        b = randdec(prec, 2 * exp_range)
+        c = randdec(prec, 2 * exp_range)
+        yield a, b, c, None
+        yield a, b, c, None, None
+
+
+class Foo:
+    CLASS_VAR = 1
+    REFERENCES_CLASS_VAR = {"CLASS_VAR": CLASS_VAR}
+
+
+class Class:
+    def __init__(self):
+        # TODO(charlie): This should be recognized as a defined variable.
+        Class  # noqa: F821
+
+
+try:
+    x = 1 / 0
+except Exception as e:
+    # TODO(charlie): This should be recognized as a defined variable.
+    print(e)  # noqa: F821

src/check_ast.rs

@@ -44,8 +44,8 @@ impl Scope {
 enum BindingKind {
     Argument,
     Assignment,
-    ClassDefinition,
     Definition,
+    ClassDefinition,
     Builtin,
     FutureImportation,
     Importation(String),
@@ -109,22 +109,22 @@ impl Visitor for Checker<'_> {
                 }
             }
             StmtKind::FunctionDef { name, .. } => {
-                self.push_scope(Scope::new(Function));
                 self.add_binding(Binding {
-                    kind: BindingKind::ClassDefinition,
+                    kind: BindingKind::Definition,
                     name: name.clone(),
                     used: None,
                     location: stmt.location,
-                })
+                });
+                self.push_scope(Scope::new(Function));
             }
             StmtKind::AsyncFunctionDef { name, .. } => {
-                self.push_scope(Scope::new(Function));
                 self.add_binding(Binding {
-                    kind: BindingKind::ClassDefinition,
+                    kind: BindingKind::Definition,
                     name: name.clone(),
                     used: None,
                     location: stmt.location,
-                })
+                });
+                self.push_scope(Scope::new(Function));
             }
             StmtKind::Return { .. } => {
                 if self
@@ -280,7 +280,7 @@ impl Visitor for Checker<'_> {
 
         if let StmtKind::ClassDef { name, .. } = &stmt.node {
             self.add_binding(Binding {
-                kind: BindingKind::Definition,
+                kind: BindingKind::ClassDefinition,
                 name: name.clone(),
                 used: None,
                 location: stmt.location,
@@ -450,11 +450,13 @@ impl Checker<'_> {
     fn handle_node_load(&mut self, expr: &Expr) {
         if let ExprKind::Name { id, .. } = &expr.node {
             let scope_id = self.scopes.last_mut().expect("No current scope found.").id;
+            let mut first_iter = true;
+            let mut in_generators = false;
             for scope in self.scopes.iter_mut().rev() {
                 if matches!(scope.kind, Class) {
                     if id == "__class__" {
                         return;
-                    } else {
+                    } else if !first_iter && !in_generators {
                         continue;
                     }
                 }
@@ -462,6 +464,9 @@ impl Checker<'_> {
                     binding.used = Some(scope_id);
                     return;
                 }
+
+                first_iter = false;
+                in_generators = matches!(scope.kind, Generator);
             }
 
             if self.settings.select.contains(&CheckCode::F821) {

src/linter.rs

@@ -275,6 +275,11 @@ mod tests {
                 location: Location::new(10, 9),
                 filename: "./resources/test/src/F821.py".to_string(),
             },
+            Message {
+                kind: CheckKind::UndefinedName("numeric_string".to_string()),
+                location: Location::new(21, 12),
+                filename: "./resources/test/src/F821.py".to_string(),
+            },
         ];
         assert_eq!(actual.len(), expected.len());
         for i in 0..actual.len() {

src/settings.rs

@@ -48,7 +48,6 @@ impl Settings {
                     CheckCode::F541,
                     CheckCode::F634,
                     CheckCode::F706,
-                    CheckCode::F821,
                     CheckCode::F831,
                     CheckCode::F832,
                     CheckCode::F901,

src/visitor.rs

@@ -340,33 +340,33 @@ pub fn walk_expr<V: Visitor + ?Sized>(visitor: &mut V, expr: &Expr) {
             }
         }
         ExprKind::ListComp { elt, generators } => {
-            visitor.visit_expr(elt);
             for comprehension in generators {
                 visitor.visit_comprehension(comprehension)
             }
+            visitor.visit_expr(elt);
         }
         ExprKind::SetComp { elt, generators } => {
-            visitor.visit_expr(elt);
             for comprehension in generators {
                 visitor.visit_comprehension(comprehension)
             }
+            visitor.visit_expr(elt);
         }
         ExprKind::DictComp {
             key,
             value,
             generators,
         } => {
+            for comprehension in generators {
+                visitor.visit_comprehension(comprehension)
+            }
             visitor.visit_expr(key);
             visitor.visit_expr(value);
-            for comprehension in generators {
-                visitor.visit_comprehension(comprehension)
-            }
         }
         ExprKind::GeneratorExp { elt, generators } => {
-            visitor.visit_expr(elt);
             for comprehension in generators {
                 visitor.visit_comprehension(comprehension)
             }
+            visitor.visit_expr(elt);
         }
         ExprKind::Await { value } => visitor.visit_expr(value),
         ExprKind::Yield { value } => {