Repository: quora/pyanalyze
Pull request #435 (2022-01-23 02:34:05)

add assert_error()


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 5c47181..1c6d81b 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,10 +2,11 @@
 
 ## Unreleased
 
-- Add support for `assert_type()` (#433)
+- Add support for `assert_error()` (#435)
+- Add support for `assert_type()` (#434)
 - `reveal_type()` and `dump_value()` now return their argument,
   the anticipated behavior for `typing.reveal_type()` in Python
-  3.11 (#432)
+  3.11 (#433)
 - Fix return type of async generator functions (#431)
 - Type check function decorators (#428)
 - Handle `NoReturn` in `async def` functions (#427)
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index ccb0541..7c438a6 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -9,6 +9,7 @@ be gracefully ignored by other type checkers.
 
 """
 from collections import defaultdict
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 import typing_extensions
 import pyanalyze
@@ -18,6 +19,7 @@ from typing import (
     Container,
     Dict,
     Iterable,
+    Iterator,
     Optional,
     Sequence,
     Tuple,
@@ -414,6 +416,22 @@ def assert_type(val: _T, typ: Any) -> _T:
     return val
 
 
+@contextmanager
+def assert_error() -> Iterator[None]:
+    """Context manager that asserts that code produces a type checker error.
+
+    Example::
+
+        with assert_error():  # ok
+            1 + "x"
+
+        with assert_error():  # error: no error found in this block
+            1 + 1
+
+    """
+    yield
+
+
 _overloads: Dict[str, List[Callable[..., Any]]] = defaultdict(list)
 _type_evaluations: Dict[str, List[Callable[..., Any]]] = defaultdict(list)
 
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 08009b8..4cfd4d1 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -69,7 +69,7 @@ from .arg_spec import ArgSpecCache, is_dot_asynq_function, UnwrapClass, IgnoredC
 from .boolability import Boolability, get_boolability
 from .checker import Checker, CheckerAttrContext
 from .error_code import ErrorCode, ERROR_DESCRIPTION
-from .extensions import ParameterTypeGuard, patch_typing_overload
+from .extensions import ParameterTypeGuard, assert_error, patch_typing_overload
 from .find_unused import UnusedObjectFinder, used
 from .options import (
     ConfigOption,
@@ -147,6 +147,7 @@ from .value import (
     AnnotatedValue,
     AnySource,
     AnyValue,
+    AssertErrorExtension,
     CanAssign,
     CanAssignError,
     ConstraintExtension,
@@ -155,6 +156,7 @@ from .value import (
     UNINITIALIZED_VALUE,
     NO_RETURN_VALUE,
     NoReturnConstraintExtension,
+    annotate_value,
     flatten_values,
     is_union,
     kv_pairs_from_mapping,
@@ -3286,16 +3288,32 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         return result
 
     def visit_With(self, node: ast.With) -> None:
-        for item in node.items:
-            self.visit_withitem(item)
+        contexts = [self.visit_withitem(item) for item in node.items]
+        if len(contexts) == 1:
+            context = contexts[0]
+            if isinstance(context, AnnotatedValue) and context.has_metadata_of_type(
+                AssertErrorExtension
+            ):
+                self._visit_assert_errors_block(node)
+                return
         self._generic_visit_list(node.body)
 
+    def _visit_assert_errors_block(self, node: ast.With) -> None:
+        with self.catch_errors() as caught:
+            self._generic_visit_list(node.body)
+        if not caught:
+            self._show_error_if_checking(
+                node,
+                "No errors found in assert_error() block",
+                error_code=ErrorCode.inference_failure,
+            )
+
     def visit_AsyncWith(self, node: ast.AsyncWith) -> None:
         for item in node.items:
             self.visit_withitem(item, is_async=True)
         self._generic_visit_list(node.body)
 
-    def visit_withitem(self, node: ast.withitem, is_async: bool = False) -> None:
+    def visit_withitem(self, node: ast.withitem, is_async: bool = False) -> Value:
         context = self.visit(node.context_expr)
         if is_async:
             protocol = "typing.AsyncContextManager"
@@ -3316,6 +3334,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         if node.optional_vars is not None:
             with qcore.override(self, "being_assigned", assigned):
                 self.visit(node.optional_vars)
+        return context
 
     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
@@ -4178,6 +4197,11 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 if caller is not None:
                     self.collector.record_call(caller, callee_val)
 
+        if (
+            isinstance(callee_wrapped, KnownValue)
+            and callee_wrapped.val is assert_error
+        ):
+            return annotate_value(return_value, [AssertErrorExtension()])
         return return_value
 
     def _can_perform_call(
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 3504a50..153f675 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -1789,6 +1789,11 @@ class AlwaysPresentExtension(Extension):
     """
 
 
+@dataclass(frozen=True)
+class AssertErrorExtension(Extension):
+    """Used for the implementation of :func:`pyanalyze.extensions.assert_error`."""
+
+
 @dataclass(frozen=True)
 class AnnotatedValue(Value):
     """Value representing a `PEP 593 <https://www.python.org/dev/peps/pep-0593/>`_ Annotated object.


Test patch:
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index 1607759..e2b5aab 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -1161,41 +1161,6 @@ class TestIssubclass(TestNameCheckVisitorBase):
                 assert_is_value(x, SubclassValue(TypedValue(int)))
 
 
-class TestInferenceHelpers(TestNameCheckVisitorBase):
-    @assert_passes()
-    def test(self) -> None:
-        from pyanalyze import dump_value, assert_is_value
-        from pyanalyze.value import Value
-
-        def capybara(val: Value) -> None:
-            reveal_type(dump_value)  # E: inference_failure
-            dump_value(reveal_type)  # E: inference_failure
-            assert_is_value(1, KnownValue(1))
-            assert_is_value(1, KnownValue(2))  # E: inference_failure
-            assert_is_value(1, val)  # E: inference_failure
-
-    @assert_passes()
-    def test_return_value(self) -> None:
-        from pyanalyze import dump_value, assert_is_value
-
-        def capybara():
-            x = dump_value(1)  # E: inference_failure
-            y = reveal_type(1)  # E: inference_failure
-            assert_is_value(x, KnownValue(1))
-            assert_is_value(y, KnownValue(1))
-
-    @assert_passes()
-    def test_assert_type(self) -> None:
-        from pyanalyze.extensions import assert_type
-        from typing import Any
-
-        def capybara(x: int) -> None:
-            assert_type(x, int)
-            assert_type(x, "int")
-            assert_type(x, Any)  # E: inference_failure
-            assert_type(x, str)  # E: inference_failure
-
-
 class TestCallableGuards(TestNameCheckVisitorBase):
     @assert_passes()
     def test_callable(self):
diff --git a/pyanalyze/test_inference_helpers.py b/pyanalyze/test_inference_helpers.py
new file mode 100644
index 0000000..ad11c16
--- /dev/null
+++ b/pyanalyze/test_inference_helpers.py
@@ -0,0 +1,55 @@
+# static analysis: ignore
+from .test_name_check_visitor import TestNameCheckVisitorBase
+from .test_node_visitor import assert_passes
+from .value import KnownValue
+
+
+class TestInferenceHelpers(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test(self) -> None:
+        from pyanalyze import dump_value, assert_is_value
+        from pyanalyze.value import Value
+
+        def capybara(val: Value) -> None:
+            reveal_type(dump_value)  # E: inference_failure
+            dump_value(reveal_type)  # E: inference_failure
+            assert_is_value(1, KnownValue(1))
+            assert_is_value(1, KnownValue(2))  # E: inference_failure
+            assert_is_value(1, val)  # E: inference_failure
+
+    @assert_passes()
+    def test_return_value(self) -> None:
+        from pyanalyze import dump_value, assert_is_value
+
+        def capybara():
+            x = dump_value(1)  # E: inference_failure
+            y = reveal_type(1)  # E: inference_failure
+            assert_is_value(x, KnownValue(1))
+            assert_is_value(y, KnownValue(1))
+
+    @assert_passes()
+    def test_assert_type(self) -> None:
+        from pyanalyze.extensions import assert_type
+        from typing import Any
+
+        def capybara(x: int) -> None:
+            assert_type(x, int)
+            assert_type(x, "int")
+            assert_type(x, Any)  # E: inference_failure
+            assert_type(x, str)  # E: inference_failure
+
+
+class TestAssertError(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test(self) -> None:
+        from pyanalyze.extensions import assert_error
+
+        def f(x: int) -> None:
+            pass
+
+        def capybara() -> None:
+            with assert_error():
+                f("x")
+
+            with assert_error():  # E: inference_failure
+                f(1)