|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
14 | 14 |
|
15 | | -import sys |
16 | | - |
17 | | - |
18 | | -def _setup_ipython_exception_handler(): |
19 | | - """Setup custom exception handler for IPython environments to ensure minimal traceback display.""" |
20 | | - try: |
21 | | - from IPython import get_ipython |
22 | | - except ImportError: |
23 | | - return |
24 | | - |
25 | | - ipython = get_ipython() |
26 | | - if ipython is None: |
27 | | - return |
28 | | - |
29 | | - # Store original method if not already stored |
30 | | - if hasattr(ipython, "_dataproc_spark_connect_original_showtraceback"): |
31 | | - return # Already patched |
32 | | - |
33 | | - ipython._dataproc_spark_connect_original_showtraceback = ( |
34 | | - ipython.showtraceback |
35 | | - ) |
36 | | - |
37 | | - def custom_showtraceback( |
38 | | - shell, |
39 | | - exc_tuple=None, |
40 | | - filename=None, |
41 | | - tb_offset=None, |
42 | | - exception_only=False, |
43 | | - running_compiled_code=False, |
44 | | - ): |
45 | | - # Get the current exception info |
46 | | - _, value, _ = sys.exc_info() if exc_tuple is None else exc_tuple |
47 | | - |
48 | | - # If it's our custom exception, show only the message |
49 | | - if isinstance(value, DataprocSparkConnectException): |
50 | | - print(f"Error: {value.message}", file=sys.stderr) |
51 | | - else: |
52 | | - # Use original behavior for other exceptions |
53 | | - shell._dataproc_spark_connect_original_showtraceback( |
54 | | - exc_tuple, |
55 | | - filename, |
56 | | - tb_offset, |
57 | | - exception_only, |
58 | | - running_compiled_code, |
59 | | - ) |
60 | | - |
61 | | - # Override the method |
62 | | - ipython.showtraceback = custom_showtraceback |
63 | | - |
64 | | - |
65 | | -# Setup the handler once at module import time |
66 | | -_setup_ipython_exception_handler() |
67 | | - |
68 | 15 |
|
69 | 16 | class DataprocSparkConnectException(Exception): |
70 | 17 | """A custom exception class to only print the error messages. |
|
0 commit comments