add catched

tests/core/test_catched.py

@@ -0,0 +1,62 @@
+from types import TracebackType
+
+import pytest
+from baby_steps import given, then, when
+from pytest import raises
+
+from vedro import catched
+
+
+def test_catched_with_expected_exception():
+    with given:
+        exc = ValueError("msg")
+
+    with when:
+        with catched(ValueError) as exc_info:
+            raise exc
+
+    with then:
+        assert exc_info.type is ValueError
+        assert exc_info.value == exc
+        assert isinstance(exc_info.traceback, TracebackType)
+
+        assert repr(exc_info) == f"<catched exception={exc!r}>"
+
+
+def test_catched_with_no_exception():
+    with when:
+        with catched(ValueError) as exc_info:
+            pass
+
+    with then:
+        assert exc_info.type is None
+        assert exc_info.value is None
+        assert exc_info.traceback is None
+
+        assert repr(exc_info) == "<catched exception=None>"
+
+
+def test_catched_with_unexpected_exception():
+    with when, raises(KeyError) as excinfo:
+        with catched(ValueError) as exc_info:
+            raise KeyError("msg")
+
+    with then:
+        assert excinfo.type is KeyError
+        assert str(excinfo.value) == "'msg'"
+
+        assert exc_info.type is None
+        assert exc_info.value is None
+        assert exc_info.traceback is None
+
+
+@pytest.mark.parametrize("exc", [ValueError("msg"), KeyError("msg")])
+def test_catched_with_multiple_exception_types(exc: Exception):
+    with when:
+        with catched((ValueError, KeyError)) as exc_info:
+            raise exc
+
+    with then:
+        assert exc_info.type is type(exc)
+        assert exc_info.value == exc
+        assert isinstance(exc_info.traceback, TracebackType)

vedro/__init__.py

@@ -3,6 +3,7 @@
 import sys
 from typing import Any
 
+from ._catched import catched
 from ._config import Config
 from ._context import context
 from ._interface import Interface
@@ -15,7 +16,7 @@
 
 __version__ = version
 __all__ = ("Scenario", "Interface", "run", "only", "skip", "skip_if", "params",
-           "context", "defer", "Config",)
+           "context", "defer", "Config", "catched",)
 
 
 def run(*, plugins: Any = None) -> None:

vedro/_catched.py

@@ -0,0 +1,44 @@
+from types import TracebackType
+from typing import Tuple, Type, Union
+
+__all__ = ("catched",)
+
+ExcpectedExcType = Union[Type[BaseException], Tuple[Type[BaseException], ...]]
+
+
+class catched:
+    def __init__(self, expected_exc: ExcpectedExcType = BaseException) -> None:
+        assert isinstance(expected_exc, tuple) or issubclass(expected_exc, BaseException)
+        self._expected_exc = expected_exc if isinstance(expected_exc, tuple) else (expected_exc,)
+        self._type: Union[Type[BaseException], None] = None
+        self._value: Union[BaseException, None] = None
+        self._traceback: Union[TracebackType, None] = None
+
+    @property
+    def type(self) -> Union[Type[BaseException], None]:
+        return self._type
+
+    @property
+    def value(self) -> Union[BaseException, None]:
+        return self._value
+
+    @property
+    def traceback(self) -> Union[TracebackType, None]:
+        return self._traceback
+
+    def __enter__(self) -> "catched":
+        return self
+
+    def __exit__(self,
+                 exc_type: Union[Type[BaseException], None],
+                 exc_value: Union[BaseException, None],
+                 traceback: Union[TracebackType, None]) -> bool:
+        if not isinstance(exc_value, self._expected_exc):
+            return False
+        self._type = exc_type
+        self._value = exc_value
+        self._traceback = traceback
+        return True
+
+    def __repr__(self) -> str:
+        return f"<catched exception={self.value!r}>"