1616import sys
1717import traceback
1818from pathlib import Path
19- from typing import Any , Callable , Dict , Optional , Union
19+ from typing import Any , Callable , Dict , List , Optional , Union
2020
2121from greenlet import greenlet
2222from pyee import AsyncIOEventEmitter
@@ -92,6 +92,32 @@ def __init__(
9292 if self ._parent :
9393 self ._parent ._objects [guid ] = self
9494
95+ def _wait_for_event_info_before (self , wait_id : str , name : str ) -> None :
96+ self ._connection ._send_message_to_server (
97+ self ._guid ,
98+ "waitForEventInfo" ,
99+ {
100+ "info" : {
101+ "name" : name ,
102+ "waitId" : wait_id ,
103+ "phase" : "before" ,
104+ "stack" : capture_call_stack (),
105+ }
106+ },
107+ )
108+
109+ def _wait_for_event_info_after (
110+ self , wait_id : str , exception : Exception = None
111+ ) -> None :
112+ info = {"waitId" : wait_id , "phase" : "after" }
113+ if exception :
114+ info ["error" ] = str (exception )
115+ self ._connection ._send_message_to_server (
116+ self ._guid ,
117+ "waitForEventInfo" ,
118+ {"info" : info },
119+ )
120+
95121 def _dispose (self ) -> None :
96122 # Clean up from parent and connection.
97123 if self ._parent :
@@ -106,7 +132,7 @@ def _dispose(self) -> None:
106132
107133class ProtocolCallback :
108134 def __init__ (self , loop : asyncio .AbstractEventLoop ) -> None :
109- self .stack_trace = "" . join ( traceback .format_stack ()[ - 10 :] )
135+ self .stack_trace : traceback . StackSummary = traceback .StackSummary ( )
110136 self .future = loop .create_future ()
111137
112138
@@ -166,14 +192,23 @@ def _send_message_to_server(
166192 ) -> ProtocolCallback :
167193 self ._last_id += 1
168194 id = self ._last_id
195+ callback = ProtocolCallback (self ._loop )
196+ if self ._is_sync :
197+ task = asyncio .current_task (self ._loop )
198+ callback .stack_trace = (
199+ getattr (task , "__pw_stack_trace__" , None ) if task else None
200+ )
201+ if not callback .stack_trace :
202+ callback .stack_trace = traceback .extract_stack ()
203+
169204 message = dict (
170205 id = id ,
171206 guid = guid ,
172207 method = method ,
173208 params = self ._replace_channels_with_guids (params , "params" ),
209+ metadata = {"stack" : serialize_call_stack (callback .stack_trace )},
174210 )
175211 self ._transport .send (message )
176- callback = ProtocolCallback (self ._loop )
177212 self ._callbacks [id ] = callback
178213 return callback
179214
@@ -184,7 +219,9 @@ def _dispatch(self, msg: ParsedMessagePayload) -> None:
184219 error = msg .get ("error" )
185220 if error :
186221 parsed_error = parse_error (error ["error" ]) # type: ignore
187- parsed_error .stack = callback .stack_trace
222+ parsed_error .stack = "" .join (
223+ traceback .format_list (callback .stack_trace )[- 10 :]
224+ )
188225 callback .future .set_exception (parsed_error )
189226 else :
190227 result = self ._replace_guids_with_channels (msg .get ("result" ))
@@ -267,3 +304,19 @@ def from_channel(channel: Channel) -> Any:
267304
268305def from_nullable_channel (channel : Optional [Channel ]) -> Optional [Any ]:
269306 return channel ._object if channel else None
307+
308+
309+ def serialize_call_stack (stack_trace : traceback .StackSummary ) -> List [Dict ]:
310+ stack : List [Dict ] = []
311+ for frame in stack_trace :
312+ if "_generated.py" in frame .filename :
313+ break
314+ stack .append (
315+ {"file" : frame .filename , "line" : frame .lineno , "function" : frame .name }
316+ )
317+ stack .reverse ()
318+ return stack
319+
320+
321+ def capture_call_stack () -> List [Dict ]:
322+ return serialize_call_stack (traceback .extract_stack ())
0 commit comments