feat: allow only_raise() to accept multiple exception types

tests/test_utils.py

@@ -3,38 +3,57 @@
 import pytest
 from tidyexc import utils, only_raise
 
-class ParentError(Exception):
+class E1(Exception):
     pass
 
-class ChildError(ParentError):
+class E1a(E1):
     pass
 
-class UnrelatedError(Exception):
+class E2(Exception):
+    pass
+
+class E3(Exception):
     pass
 
 @pytest.mark.parametrize(
         'only_err, raise_err, catch_err', [
-            (   ParentError,    ParentError,    ParentError),
-            (   ParentError,     ChildError,     ChildError),
-            (   ParentError, UnrelatedError,    ParentError),
+            ([E1], E1,  E1),
+            ([E1], E1a, E1a),
+            ([E1], E2,  E1),
+
+            ([E1a], E1,  E1a),
+            ([E1a], E1a, E1a),
+            ([E1a], E2,  E1a),
 
-            (    ChildError,    ParentError,     ChildError),
-            (    ChildError,     ChildError,     ChildError),
-            (    ChildError, UnrelatedError,     ChildError),
+            ([E2], E1,  E2),
+            ([E2], E1a, E2),
+            ([E2], E2,  E2),
 
-            (UnrelatedError,    ParentError, UnrelatedError),
-            (UnrelatedError,     ChildError, UnrelatedError),
-            (UnrelatedError, UnrelatedError, UnrelatedError),
+            ([E1, E2], E1,  E1),
+            ([E1, E2], E1a, E1a),
+            ([E1, E2], E2,  E2),
+            ([E1, E2], E3,  E1),
+
+            ([E2, E1], E1,  E1),
+            ([E2, E1], E1a, E1a),
+            ([E2, E1], E2,  E2),
+            ([E2, E1], E3,  E2),
 ])
 def test_only_raise(only_err, raise_err, catch_err):
 
-    @only_raise(only_err)
+    @only_raise(*only_err)
     def f(err):
         raise err
 
     with pytest.raises(catch_err):
         f(raise_err)
 
+def test_only_raise_no_args_err():
+    with pytest.raises(TypeError):
+        @only_raise()
+        def f():
+            pass
+
 def test_list_iadd():
     l = utils.list_iadd()
 

tidyexc/utils.py

@@ -5,15 +5,18 @@
 
 class only_raise:
     """
-    Guarantee that the decorated function can only raise the given type of 
+    Guarantee that the decorated function can only raise the given types of 
     exception.
 
-    Any unhandled exception raised by the decorated function will be caught and 
-    re-raised using an exception of the given type. 
+    At least one type of exceptions must be specified.  Any unhandled exception 
+    raised by the decorated function will be caught and re-raised using an 
+    exception of the first given type.
     """
 
-    def __init__(self, err_cls):
-        self.err_cls = err_cls
+    def __init__(self, *err_cls):
+        if not err_cls:
+            raise TypeError("__init__() missing 1 required positional argument: 'err_cls'")
+        self.err_cls = tuple(err_cls)
 
     def __call__(self, f):
 
@@ -24,7 +27,7 @@ def wrapper(*args, **kwargs):
             except self.err_cls as err:
                 raise err from None
             except Exception as err:
-                raise self.err_cls(str(err)) from err
+                raise self.err_cls[0](str(err)) from err
 
         return wrapper