Return self from context managers

Author: srittau

NEWS.md

@@ -6,6 +6,19 @@ API-Incompatible Changes
 
 * Drop support for Python 3.4.
 
+API Additions
+-------------
+
+* `AssertRaisesContext` and `AssertWarnsContext` now return themselves
+  when `__enter__()` is called. By extension it now easier to call
+  `add_test()` with `assert_raises()` et al:
+
+```python
+with assert_raises(KeyError) as context:
+    context.add_test(...)
+    ...
+```  
+
 News in asserts 0.9.1
 =====================
 

asserts/__init__.py

@@ -884,7 +884,7 @@ def __init__(self, exception, msg_fmt="{msg}"):
         self._tests = []
 
     def __enter__(self):
-        pass
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if not exc_type:
@@ -1166,6 +1166,7 @@ def __init__(self, warning_class, msg_fmt="{msg}"):
     def __enter__(self):
         self._warning_context = catch_warnings(record=True)
         self._warnings = self._warning_context.__enter__()
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self._warning_context.__exit__(exc_type, exc_val, exc_tb)

asserts/__init__.pyi

@@ -19,12 +19,13 @@ from typing import (
 )
 
 _E = TypeVar("_E", bound=BaseException)
+_S = TypeVar("_S")
 
 class AssertRaisesContext(Generic[_E]):
     exception: Type[_E]
     msg_fmt: Text
     def __init__(self, exception: Type[_E], msg_fmt: Text = ...) -> None: ...
-    def __enter__(self) -> AssertRaisesContext: ...
+    def __enter__(self: _S) -> _S: ...
     def __exit__(
         self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
     ) -> Optional[bool]: ...
@@ -41,7 +42,7 @@ class AssertRaisesRegexContext(AssertRaisesContext[_E]):
 
 class AssertWarnsContext:
     def __init__(self, warning_class: Type[Warning], msg_fmt: Text = ...) -> None: ...
-    def __enter__(self) -> None: ...
+    def __enter__(self: _S) -> _S: ...
     def __exit__(self, exc_type: Type[BaseException], exc_val: BaseException, exc_tb: Any) -> None: ...
     def format_message(self) -> Text: ...
     def add_test(self, cb: Callable[[Warning], None]) -> None: ...

test_asserts.py

@@ -47,6 +47,11 @@
 )
 
 
+class Box:
+    def __init__(self, initial_value):
+        self.value = initial_value
+
+
 class _DummyObject(object):
     def __init__(self, value="x"):
         self.value = value
@@ -962,6 +967,29 @@ def test_assert_raises__wrong_exception_raised(self):
         else:
             fail("no exception raised")
 
+    def test_assert_raises__add_test_called(self):
+        called = Box(False)
+
+        def extra_test(exc):
+            assert_is_instance(exc, KeyError)
+            called.value = True
+
+        with assert_raises(KeyError) as context:
+            context.add_test(extra_test)
+            raise KeyError()
+        assert_true(called.value, "extra_test() was not called")
+
+    def test_assert_raises__add_test_not_called(self):
+        called = Box(False)
+
+        def extra_test(_):
+            called.value = True
+
+        with assert_raises(AssertionError):
+            with assert_raises(KeyError) as context:
+                context.add_test(extra_test)
+        assert_false(called.value, "extra_test() was unexpectedly called")
+
     # assert_raises_regex()
 
     def test_assert_raises_regex__raises_right_exception(self):
@@ -1166,6 +1194,30 @@ def test_assert_warns__warning_handler_deinstalled_on_failure(self):
             warn("bar", UserWarning)
             assert_equal(1, len(warnings))
 
+    def test_assert_warns__add_test_called(self):
+        called = Box(False)
+
+        def extra_test(warning):
+            assert_is(warning.category, UserWarning)
+            called.value = True
+            return True
+
+        with assert_warns(UserWarning) as context:
+            context.add_test(extra_test)
+            warn("bar", UserWarning)
+        assert_true(called.value, "extra_test() was not called")
+
+    def test_assert_warns__add_test_not_called(self):
+        called = Box(False)
+
+        def extra_test(_):
+            called.value = True
+
+        with assert_raises(AssertionError):
+            with assert_warns(UserWarning) as context:
+                context.add_test(extra_test)
+        assert_false(called.value, "extra_test() was unexpectedly called")
+
     # assert_warns_regex()
 
     def test_assert_warns_regex__warned(self):