-
-
Notifications
You must be signed in to change notification settings - Fork 548
/
Copy pathhandler.py
99 lines (70 loc) · 2.86 KB
/
handler.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import os
import sys
import threading
from types import TracebackType
from typing import Any, Callable, Optional, Tuple, Type, cast
from .exception import StrawberryException, UnableToFindExceptionSource
original_exception_hook = sys.excepthook
if sys.version_info >= (3, 8):
original_threading_exception_hook = threading.excepthook
else:
original_threading_exception_hook = None
ExceptionHandler = Callable[
[Type[BaseException], BaseException, Optional[TracebackType]], None
]
def should_use_rich_exceptions():
errors_disabled = os.environ.get("STRAWBERRY_DISABLE_RICH_ERRORS", "")
return errors_disabled.lower() not in ["true", "1", "yes"]
def _get_handler(exception_type: Type[BaseException]) -> ExceptionHandler:
if issubclass(exception_type, StrawberryException):
try:
import rich
except ImportError:
pass
else:
def _handler(
exception_type: Type[BaseException],
exception: BaseException,
traceback: Optional[TracebackType],
):
try:
rich.print(exception)
# we check if weren't able to find the exception source
# in that case we fallback to the original exception handler
except UnableToFindExceptionSource:
original_exception_hook(exception_type, exception, traceback)
return _handler
return original_exception_hook
def strawberry_exception_handler(
exception_type: Type[BaseException],
exception: BaseException,
traceback: Optional[TracebackType],
):
_get_handler(exception_type)(exception_type, exception, traceback)
def strawberry_threading_exception_handler(
args: Tuple[
Type[BaseException],
Optional[BaseException],
Optional[TracebackType],
Optional[threading.Thread],
]
):
(exception_type, exception, traceback, _) = args
if exception is None:
if sys.version_info >= (3, 8):
# this cast is only here because some weird issue with mypy
# and the inability to disable this error based on the python version
# (we'd need to do type ignore for python 3.8 and above, but mypy
# doesn't seem to be able to handle that and will complain in python 3.7)
cast(Any, original_threading_exception_hook)(args)
return
_get_handler(exception_type)(exception_type, exception, traceback)
def reset_exception_handler():
sys.excepthook = original_exception_hook
if sys.version_info >= (3, 8):
threading.excepthook = original_threading_exception_hook
def setup_exception_handler():
if should_use_rich_exceptions():
sys.excepthook = strawberry_exception_handler
if sys.version_info >= (3, 8):
threading.excepthook = strawberry_threading_exception_handler