1212# See the License for the specific language governing permissions and
1313# limitations under the License.
1414
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+
1568
1669class DataprocSparkConnectException (Exception ):
1770 """A custom exception class to only print the error messages.
@@ -24,4 +77,4 @@ def __init__(self, message):
2477 super ().__init__ (message )
2578
2679 def _render_traceback_ (self ):
27- return self .message
80+ return [ self .message ]
0 commit comments