Avoid F821 false-positives with NameError (#386)

resources/test/fixtures/F821.py

@@ -122,3 +122,12 @@ class PEP593Test:
         dict["foo", "bar"],  # Expected to fail as undefined.
         123,
     ]
+
+
+def in_ipython_notebook() -> bool:
+    try:
+        # autoimported by notebooks
+        get_ipython()  # type: ignore[name-defined]
+    except NameError:
+        return False  # not in notebook
+    return True

src/ast/helpers.rs

@@ -1,6 +1,6 @@
 use once_cell::sync::Lazy;
 use regex::Regex;
-use rustpython_ast::{Expr, ExprKind, StmtKind};
+use rustpython_ast::{Excepthandler, ExcepthandlerKind, Expr, ExprKind, StmtKind};
 
 use crate::python::typing;
 
@@ -35,6 +35,22 @@ pub fn match_annotated_subscript(expr: &Expr) -> Option<SubscriptKind> {
     }
 }
 
+fn node_name(expr: &Expr) -> Option<&str> {
+    if let ExprKind::Name { id, .. } = &expr.node {
+        Some(id)
+    } else {
+        None
+    }
+}
+
+pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
+    match &expr.node {
+        ExprKind::Attribute { attr, .. } => target == attr,
+        ExprKind::Name { id, .. } => target == id,
+        _ => false,
+    }
+}
+
 pub fn is_assignment_to_a_dunder(node: &StmtKind) -> bool {
     // Check whether it's an assignment to a dunder, with or without a type annotation.
     // This is what pycodestyle (as of 2.9.1) does.
@@ -65,10 +81,27 @@ pub fn is_assignment_to_a_dunder(node: &StmtKind) -> bool {
     }
 }
 
-pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
-    match &expr.node {
-        ExprKind::Attribute { attr, .. } => target == attr,
-        ExprKind::Name { id, .. } => target == id,
-        _ => false,
+/// Extract the names of all handled exceptions.
+/// Note that, for now, this only matches on ExprKind::Name, and so won't catch exceptions like
+/// `module.CustomException`. (But will catch all builtin exceptions.)
+pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec<&str> {
+    let mut handler_names = vec![];
+    for handler in handlers {
+        match &handler.node {
+            ExcepthandlerKind::ExceptHandler { type_, .. } => {
+                if let Some(type_) = type_ {
+                    if let ExprKind::Tuple { elts, .. } = &type_.node {
+                        for type_ in elts {
+                            if let Some(name) = node_name(type_) {
+                                handler_names.push(name);
+                            }
+                        }
+                    } else if let Some(name) = node_name(type_) {
+                        handler_names.push(name);
+                    }
+                }
+            }
+        }
     }
+    handler_names
 }

src/check_ast.rs

@@ -10,7 +10,7 @@ use rustpython_parser::ast::{
 };
 use rustpython_parser::parser;
 
-use crate::ast::helpers::{match_name_or_attr, SubscriptKind};
+use crate::ast::helpers::{extract_handler_names, match_name_or_attr, SubscriptKind};
 use crate::ast::operations::{extract_all_names, SourceCodeLocator};
 use crate::ast::relocate::relocate_expr;
 use crate::ast::types::{
@@ -60,6 +60,7 @@ pub struct Checker<'a> {
     seen_docstring: bool,
     futures_allowed: bool,
     annotations_future_enabled: bool,
+    except_handlers: Vec<Vec<&'a str>>,
 }
 
 impl<'a> Checker<'a> {
@@ -74,25 +75,26 @@ impl<'a> Checker<'a> {
             autofix,
             path,
             locator: SourceCodeLocator::new(content),
-            checks: vec![],
-            parents: vec![],
-            parent_stack: vec![],
-            scopes: vec![],
-            scope_stack: vec![],
-            dead_scopes: vec![],
-            deferred_string_annotations: vec![],
-            deferred_annotations: vec![],
-            deferred_functions: vec![],
-            deferred_lambdas: vec![],
-            deferred_assignments: vec![],
-            in_f_string: None,
-            in_annotation: false,
-            in_literal: false,
-            seen_non_import: false,
-            seen_docstring: false,
-            futures_allowed: true,
-            annotations_future_enabled: false,
+            checks: Default::default(),
             deletions: Default::default(),
+            parents: Default::default(),
+            parent_stack: Default::default(),
+            scopes: Default::default(),
+            scope_stack: Default::default(),
+            dead_scopes: Default::default(),
+            deferred_string_annotations: Default::default(),
+            deferred_annotations: Default::default(),
+            deferred_functions: Default::default(),
+            deferred_lambdas: Default::default(),
+            deferred_assignments: Default::default(),
+            in_f_string: None,
+            in_annotation: Default::default(),
+            in_literal: Default::default(),
+            seen_non_import: Default::default(),
+            seen_docstring: Default::default(),
+            futures_allowed: true,
+            annotations_future_enabled: Default::default(),
+            except_handlers: Default::default(),
         }
     }
 }
@@ -601,6 +603,27 @@ where
                     self.visit_stmt(stmt);
                 }
             }
+            StmtKind::Try {
+                body,
+                handlers,
+                orelse,
+                finalbody,
+            } => {
+                self.except_handlers.push(extract_handler_names(handlers));
+                for stmt in body {
+                    self.visit_stmt(stmt);
+                }
+                self.except_handlers.pop();
+                for excepthandler in handlers {
+                    self.visit_excepthandler(excepthandler)
+                }
+                for stmt in orelse {
+                    self.visit_stmt(stmt);
+                }
+                for stmt in finalbody {
+                    self.visit_stmt(stmt);
+                }
+            }
             _ => visitor::walk_stmt(self, stmt),
         };
 
@@ -1486,6 +1509,14 @@ impl<'a> Checker<'a> {
                 if self.path.ends_with("__init__.py") && id == "__path__" {
                     return;
                 }
+
+                // Avoid flagging if NameError is handled.
+                if let Some(handler_names) = self.except_handlers.last() {
+                    if handler_names.contains(&"NameError") {
+                        return;
+                    }
+                }
+
                 self.checks.push(Check::new(
                     CheckKind::UndefinedName(id.clone()),
                     self.locate_check(Range::from_located(expr)),