Merge pull request #881 from jvansanten/count-multiple-exceptions

Allow tuples of exceptions in ExceptionCounter

Author: csmarchbanks

prometheus_client/context_managers.py

@@ -1,7 +1,9 @@
 import sys
 from timeit import default_timer
 from types import TracebackType
-from typing import Any, Callable, Optional, Type, TYPE_CHECKING, TypeVar
+from typing import (
+    Any, Callable, Optional, Tuple, Type, TYPE_CHECKING, TypeVar, Union,
+)
 
 if sys.version_info >= (3, 8, 0):
     from typing import Literal
@@ -14,7 +16,7 @@
 
 
 class ExceptionCounter:
-    def __init__(self, counter: "Counter", exception: Type[BaseException]) -> None:
+    def __init__(self, counter: "Counter", exception: Union[Type[BaseException], Tuple[Type[BaseException], ...]]) -> None:
         self._counter = counter
         self._exception = exception
 

prometheus_client/metrics.py

@@ -3,8 +3,8 @@
 import time
 import types
 from typing import (
-    Any, Callable, Dict, Iterable, List, Optional, Sequence, Type, TypeVar,
-    Union,
+    Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Type,
+    TypeVar, Union,
 )
 
 from . import values  # retain this import style for testability
@@ -288,7 +288,7 @@ def inc(self, amount: float = 1, exemplar: Optional[Dict[str, str]] = None) -> N
             _validate_exemplar(exemplar)
             self._value.set_exemplar(Exemplar(exemplar, amount, time.time()))
 
-    def count_exceptions(self, exception: Type[BaseException] = Exception) -> ExceptionCounter:
+    def count_exceptions(self, exception: Union[Type[BaseException], Tuple[Type[BaseException], ...]] = Exception) -> ExceptionCounter:
         """Count exceptions in a block of code or function.
 
         Can be used as a function decorator or context manager.