Repository: quora/pyanalyze
Pull request #504 (2022-03-19 20:43:23)

Fix scope tracking inside  try and with blocks


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 28388f1..49c3d18 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Correctly track assignments of variables inside `try` blocks
+  and inside `with` blocks that may suppress exceptions (#504)
 - Support mappings that do not inherit from `collections.abc.Mapping`
   (#501)
 - Improve type inference for calls to `set()`, `list()`, and
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index b090dd4..2b83a57 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -3473,17 +3473,11 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             return
         first_item = items[0]
         can_suppress = self.visit_withitem(first_item, is_async)
-        with self.scopes.subscope() as rest_scope:
-            # get scope for visiting remaining CMs + Body
-            self.visit_single_cm(items[1:], body, is_async=is_async)
         if can_suppress:
-            # If an exception was suppressed, assume no other CMs
-            # or any code in the body was executed.
-            with self.scopes.subscope() as dummy_subscope:
-                pass
-            self.scopes.combine_subscopes([rest_scope, dummy_subscope])
+            with self.scopes.suppressing_subscope():
+                self.visit_single_cm(items[1:], body, is_async=is_async)
         else:
-            self.scopes.combine_subscopes([rest_scope])
+            self.visit_single_cm(items[1:], body, is_async=is_async)
 
     def visit_withitem(self, node: ast.withitem, is_async: bool = False) -> bool:
         context = self.visit(node.context_expr)
@@ -3520,48 +3514,50 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 self.visit(node.optional_vars)
         return can_suppress
 
-    def visit_try_except(self, node: ast.Try) -> List[SubScope]:
-        # reset yield checks between branches to avoid incorrect errors when we yield both in the
-        # try and the except block
+    def visit_try_except(self, node: ast.Try) -> None:
         with self.scopes.subscope():
-            with self.scopes.subscope() as try_scope:
-                self._generic_visit_list(node.body)
+            with self.scopes.subscope() as dummy_scope:
+                pass
+
+            with self.scopes.subscope() as failure_scope:
+                with self.scopes.suppressing_subscope() as success_scope:
+                    self._generic_visit_list(node.body)
+
+            with self.scopes.subscope() as else_scope:
                 self.yield_checker.reset_yield_checks()
+                self.scopes.combine_subscopes([success_scope])
                 self._generic_visit_list(node.orelse)
-            with self.scopes.subscope() as dummy_subscope:
-                pass
-            self.scopes.combine_subscopes([try_scope, dummy_subscope])
 
             except_scopes = []
             for handler in node.handlers:
                 with self.scopes.subscope() as except_scope:
                     except_scopes.append(except_scope)
+                    # reset yield checks between branches to avoid incorrect errors when we yield
+                    # both in the try and the except block
                     self.yield_checker.reset_yield_checks()
+                    self.scopes.combine_subscopes([dummy_scope, failure_scope])
                     self.visit(handler)
 
-        return [try_scope] + except_scopes
+        self.scopes.combine_subscopes([else_scope, *except_scopes])
 
     def visit_Try(self, node: ast.Try) -> None:
         # py3 combines the Try and Try/Finally nodes
         if node.finalbody:
-            subscopes = self.visit_try_except(node)
+            with self.scopes.subscope() as failure_scope:
+                with self.scopes.suppressing_subscope() as success_scope:
+                    self.visit_try_except(node)
 
-            # For the case where nothing in the try-except block is executed
+            # If the try block fails
             with self.scopes.subscope():
-                self._generic_visit_list(node.finalbody)
-
-            # For the case where something in the try-except exits the scope
-            with self.scopes.subscope():
-                self.scopes.combine_subscopes(subscopes, ignore_leaves_scope=True)
+                self.scopes.combine_subscopes([failure_scope])
                 self._generic_visit_list(node.finalbody)
 
             # For the case where execution continues after the try-finally
-            self.scopes.combine_subscopes(subscopes)
+            self.scopes.combine_subscopes([success_scope])
             self._generic_visit_list(node.finalbody)
         else:
             # Life is much simpler without finally
-            subscopes = self.visit_try_except(node)
-            self.scopes.combine_subscopes(subscopes)
+            self.visit_try_except(node)
         self.yield_checker.reset_yield_checks()
 
     def visit_ExceptHandler(self, node: ast.ExceptHandler) -> None:
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index 76b39c7..bd39239 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -797,6 +797,10 @@ class Scope:
     def __contains__(self, varname: Varname) -> bool:
         return varname in self.variables
 
+    @contextlib.contextmanager
+    def suppressing_subscope(self) -> Iterator[SubScope]:
+        yield {}
+
     # no real subscopes in non-function scopes, just dummy implementations
     @contextlib.contextmanager
     def subscope(self) -> Iterator[None]:
@@ -1156,6 +1160,53 @@ class FunctionScope(Scope):
                 return EMPTY_ORIGIN
         return self._resolve_origin(definers)
 
+    def get_all_definition_nodes(self) -> Dict[Varname, Set[Node]]:
+        """Return a copy of name_to_all_definition_nodes."""
+        return {
+            key: set(nodes) for key, nodes in self.name_to_all_definition_nodes.items()
+        }
+
+    @contextlib.contextmanager
+    def suppressing_subscope(self) -> Iterator[SubScope]:
+        """A suppressing subscope is a subscope that may suppress exceptions
+        inside of it.
+
+        This is used to implement try and with blocks. After code like this::
+
+            x = 1
+            try:
+                x = 2
+                x = 3
+            except Exception:
+                pass
+
+        The value of `x` may be any of 1, 2, and 3, depending on whether and
+        where an exception was thrown.
+
+        To implement this, we keep track of all assignments inside the block
+        and give them effect, so that after the suppressing subscope ends,
+        each variable's definition nodes include all of these assignments.
+
+        """
+        old_defn_nodes = self.get_all_definition_nodes()
+        with self.subscope() as inner_scope:
+            yield inner_scope
+        new_defn_nodes = self.get_all_definition_nodes()
+        rest_scope = {
+            key: list(nodes - old_defn_nodes.get(key, set()))
+            for key, nodes in new_defn_nodes.items()
+            if key != LEAVES_SCOPE
+        }
+        rest_scope = {key: nodes for key, nodes in rest_scope.items() if nodes}
+        with self.subscope() as dummy_subscope:
+            pass
+        all_keys = set(rest_scope) | set(dummy_subscope)
+        new_scope = {
+            key: [*dummy_subscope.get(key, []), *rest_scope.get(key, [])]
+            for key in all_keys
+        }
+        self.combine_subscopes([dummy_subscope, new_scope])
+
     @contextlib.contextmanager
     def subscope(self) -> Iterator[SubScope]:
         """Create a new subscope, to be used for conditional branches."""
@@ -1430,6 +1481,9 @@ class StackedScopes:
         """
         self.scopes[-1].set(varname, value, node, state)
 
+    def suppressing_subscope(self) -> ContextManager[SubScope]:
+        return self.scopes[-1].suppressing_subscope()
+
     def subscope(self) -> ContextManager[SubScope]:
         """Creates a new subscope (see the :class:`FunctionScope` docstring)."""
         return self.scopes[-1].subscope()


Test patch:
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index 95678b6..e1c4f87 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -2100,26 +2100,34 @@ class TestContextManagerWithSuppression(TestNameCheckVisitorBase):
                 a = 3
             assert_is_value(a, KnownValue(2) | KnownValue(3))
 
+        def use_suppress_exception_multi_assignment():
+            a = 2
+            with SuppressException():
+                a = 3
+                a = 4
+            assert_is_value(a, KnownValue(2) | KnownValue(3) | KnownValue(4))
+
         def use_empty_context():
-            a = 2  # static analysis: ignore[unused_variable]
+            a = 2  # E: unused_variable
             with EmptyContext():
-                a = 3
-            assert_is_value(a, KnownValue(3))
+                a = 3  # E: unused_variable
+                a = 4
+            assert_is_value(a, KnownValue(4))
 
         def use_context_manager():
-            a = 2  # static analysis: ignore[unused_variable]
+            a = 2  # E: unused_variable
             with empty_context_manager():
                 a = 3
             assert_is_value(a, KnownValue(3))
 
         def use_builtin_function():
-            a = 2  # static analysis: ignore[unused_variable]
+            a = 2  # E: unused_variable
             with open("test_file.txt"):
                 a = 3
             assert_is_value(a, KnownValue(3))
 
         def use_contextlib_manager():
-            a = 2  # static analysis: ignore[unused_variable]
+            a = 2  # E: unused_variable
             with empty_contextlib_manager():
                 a = 3
             assert_is_value(a, KnownValue(3))
@@ -2130,11 +2138,38 @@ class TestContextManagerWithSuppression(TestNameCheckVisitorBase):
                 assert_is_value(b, KnownValue(None))
             assert_is_value(b, KnownValue(2) | KnownValue(None))
 
-            c = 2  # static analysis: ignore[unused_variable]
+            c = 2  # E: unused_variable
             with EmptyContext() as c, SuppressException():
                 assert_is_value(c, KnownValue(None))
             assert_is_value(c, KnownValue(None))
 
+    @assert_passes()
+    def test_possibly_undefined_with_leaves_scope(self):
+        from typing import Optional, Type
+        from types import TracebackType
+
+        class SuppressException:
+            def __enter__(self) -> None:
+                pass
+
+            def __exit__(
+                self,
+                typ: Optional[Type[BaseException]],
+                exn: Optional[BaseException],
+                tb: Optional[TracebackType],
+            ) -> bool:
+                return isinstance(exn, Exception)
+
+        def use_suppress_with_nested_block():
+            with SuppressException():
+                a = 4
+                try:
+                    b = 3
+                except Exception:
+                    return
+            print(a)  # E: possibly_undefined_name
+            print(b)  # E: possibly_undefined_name
+
     @assert_passes()
     def test_async(self):
         from typing import Optional, Type, AsyncContextManager
@@ -2174,13 +2209,13 @@ class TestContextManagerWithSuppression(TestNameCheckVisitorBase):
             assert_is_value(a, KnownValue(2) | KnownValue(3))
 
         async def use_async_empty_context():
-            a = 2  # static analysis: ignore[unused_variable]
+            a = 2  # E: unused_variable
             async with AsyncEmptyContext():
                 a = 3
             assert_is_value(a, KnownValue(3))
 
         async def use_async_context_manager():
-            a = 2  # static analysis: ignore[unused_variable]
+            a = 2  # E: unused_variable
             async with async_empty_context_manager():
                 a = 3
             assert_is_value(a, KnownValue(3))
@@ -2191,7 +2226,7 @@ class TestContextManagerWithSuppression(TestNameCheckVisitorBase):
                 assert_is_value(b, KnownValue(None))
             assert_is_value(b, KnownValue(2) | KnownValue(None))
 
-            c = 2  # static analysis: ignore[unused_variable]
+            c = 2  # E: unused_variable
             async with AsyncEmptyContext() as c, AsyncSuppressException():
                 assert_is_value(c, KnownValue(None))
             assert_is_value(c, KnownValue(None))
@@ -2206,7 +2241,7 @@ class TestContextManagerWithSuppression(TestNameCheckVisitorBase):
             yield
 
         async def use_async_contextlib_manager():
-            a = 2  # static analysis: ignore[unused_variable]
+            a = 2  # E: unused_variable
             async with async_empty_contextlib_manager():
                 a = 3
             assert_is_value(a, KnownValue(3))
diff --git a/pyanalyze/test_stacked_scopes.py b/pyanalyze/test_stacked_scopes.py
index 7a8ffee..64a2180 100644
--- a/pyanalyze/test_stacked_scopes.py
+++ b/pyanalyze/test_stacked_scopes.py
@@ -310,16 +310,34 @@ class TestTry(TestNameCheckVisitorBase):
     @assert_passes()
     def test_else(self):
         def capybara():
+            x = 1
             try:
+                x = 2
                 x = 3
                 assert_is_value(x, KnownValue(3))
             except NameError:
+                assert_is_value(x, KnownValue(1) | KnownValue(2) | KnownValue(3))
                 x = 4
                 assert_is_value(x, KnownValue(4))
             else:
+                assert_is_value(x, KnownValue(3))
                 x = 5
                 assert_is_value(x, KnownValue(5))
-            assert_is_value(x, MultiValuedValue([KnownValue(5), KnownValue(4)]))
+            assert_is_value(x, KnownValue(5) | KnownValue(4))
+
+    @assert_passes()
+    def test_multiple_assignment(self):
+        def capybara():
+            x = 1
+            try:
+                x = 2
+                x = 3
+                assert_is_value(x, KnownValue(3))
+            except NameError:
+                assert_is_value(x, KnownValue(1) | KnownValue(2) | KnownValue(3))
+            else:
+                assert_is_value(x, KnownValue(3))
+            assert_is_value(x, KnownValue(1) | KnownValue(2) | KnownValue(3))
 
     @assert_passes()
     def test_finally(self):