Fix globals scanning for exception handlers

Author: jonatanklosko

lib/pythonx/ast/scan_globals.py

@@ -9,15 +9,25 @@ def __init__(self):
     self.scopes = []
     # Add the global scope to the stack.
     self._push_scope()
+    # Scope information specific to exception handlers.
+    self.exception_handler_scopes = []
 
   def _push_scope(self, is_comprehension=False):
     self.scopes.append({"defs": set(), "is_comprehension": is_comprehension})
 
   def _pop_scope(self):
     self.scopes.pop()
 
+  def _push_exception_handler_scope(self):
+    self.exception_handler_scopes.append({"defs": set()})
+
+  def _pop_exception_handler_scope(self):
+    self.exception_handler_scopes.pop()
+
   def _handle_ref(self, name):
-    is_defined = any(name in scope["defs"] for scope in self.scopes)
+    is_defined = any(name in scope["defs"] for scope in self.scopes) or any(
+      name in scope["defs"] for scope in self.exception_handler_scopes
+    )
 
     if not is_defined:
       self.refs.add(name)
@@ -164,6 +174,15 @@ def visit_AnnAssign(self, node):
 
     self.visit(node.target)
 
+  def visit_ExceptHandler(self, node):
+    # Exception handlers define local variables available only within
+    # the handler, however the handler body is still within the outer
+    # scope, so we add a special ref-only scope.
+    self._push_exception_handler_scope()
+    self.exception_handler_scopes[-1]["defs"].add(node.name)
+    self.generic_visit(node)
+    self._pop_exception_handler_scope()
+
   def visit_MatchAs(self, node):
     if node.name is not None:
       self.scopes[-1]["defs"].add(node.name)

test/pythonx/ast_test.exs

@@ -200,6 +200,25 @@ defmodule Pythonx.ASTTest do
              }
     end
 
+    test "try" do
+      assert Pythonx.AST.scan_globals("""
+             try:
+               raise RuntimeError("error")
+             except RuntimeError as e:
+               x = e
+             """) == %{referenced: MapSet.new([]), defined: MapSet.new(["x"])}
+
+      assert Pythonx.AST.scan_globals("""
+             try:
+               raise RuntimeError("error")
+             except RuntimeError as e1:
+               try:
+                 raise RuntimeError("error")
+               except RuntimeError as e2:
+                 x = (e1, e2)
+             """) == %{referenced: MapSet.new([]), defined: MapSet.new(["x"])}
+    end
+
     test "global in top-level expression" do
       assert Pythonx.AST.scan_globals("""
              x + 1