1414
1515import asyncio
1616import contextvars
17+ import datetime
1718import inspect
1819import sys
1920import traceback
3334 from playwright ._impl ._playwright import Playwright
3435
3536
37+ if sys .version_info >= (3 , 8 ): # pragma: no cover
38+ from typing import TypedDict
39+ else : # pragma: no cover
40+ from typing_extensions import TypedDict
41+
42+
3643class Channel (AsyncIOEventEmitter ):
3744 def __init__ (self , connection : "Connection" , guid : str ) -> None :
3845 super ().__init__ ()
@@ -220,10 +227,11 @@ def __init__(
220227 self ._error : Optional [BaseException ] = None
221228 self .is_remote = False
222229 self ._init_task : Optional [asyncio .Task ] = None
223- self ._api_zone : contextvars .ContextVar [Optional [ Dict ]] = contextvars . ContextVar (
224- "ApiZone" , default = None
225- )
230+ self ._api_zone : contextvars .ContextVar [
231+ Optional [ ParsedStackTrace ]
232+ ] = contextvars . ContextVar ( "ApiZone" , default = None )
226233 self ._local_utils : Optional ["LocalUtils" ] = local_utils
234+ self ._stack_collector : List [List [Dict [str , Any ]]] = []
227235
228236 @property
229237 def local_utils (self ) -> "LocalUtils" :
@@ -271,6 +279,13 @@ def call_on_object_with_known_name(
271279 ) -> None :
272280 self ._waiting_for_object [guid ] = callback
273281
282+ def start_collecting_call_metadata (self , collector : Any ) -> None :
283+ if collector not in self ._stack_collector :
284+ self ._stack_collector .append (collector )
285+
286+ def stop_collecting_call_metadata (self , collector : Any ) -> None :
287+ self ._stack_collector .remove (collector )
288+
274289 def _send_message_to_server (
275290 self , guid : str , method : str , params : Dict
276291 ) -> ProtocolCallback :
@@ -283,12 +298,30 @@ def _send_message_to_server(
283298 getattr (task , "__pw_stack_trace__" , traceback .extract_stack ()),
284299 )
285300 self ._callbacks [id ] = callback
301+ stack_trace_information = cast (ParsedStackTrace , self ._api_zone .get ())
302+ for collector in self ._stack_collector :
303+ collector .append ({"stack" : stack_trace_information ["frames" ], "id" : id })
304+ frames = stack_trace_information .get ("frames" , [])
305+ location = (
306+ {
307+ "file" : frames [0 ]["file" ],
308+ "line" : frames [0 ]["line" ],
309+ "column" : frames [0 ]["column" ],
310+ }
311+ if len (frames ) > 0
312+ else None
313+ )
286314 message = {
287315 "id" : id ,
288316 "guid" : guid ,
289317 "method" : method ,
290318 "params" : self ._replace_channels_with_guids (params ),
291- "metadata" : self ._api_zone .get (),
319+ "metadata" : {
320+ "wallTime" : int (datetime .datetime .now ().timestamp () * 1000 ),
321+ "apiName" : stack_trace_information ["apiName" ],
322+ "location" : location ,
323+ "internal" : not stack_trace_information ["apiName" ],
324+ },
292325 }
293326 self ._transport .send (message )
294327 self ._callbacks [id ] = callback
@@ -412,9 +445,7 @@ async def wrap_api_call(
412445 return await cb ()
413446 task = asyncio .current_task (self ._loop )
414447 st : List [inspect .FrameInfo ] = getattr (task , "__pw_stack__" , inspect .stack ())
415- metadata = _extract_metadata_from_stack (st , is_internal )
416- if metadata :
417- self ._api_zone .set (metadata )
448+ self ._api_zone .set (_extract_stack_trace_information_from_stack (st , is_internal ))
418449 try :
419450 return await cb ()
420451 finally :
@@ -427,9 +458,7 @@ def wrap_api_call_sync(
427458 return cb ()
428459 task = asyncio .current_task (self ._loop )
429460 st : List [inspect .FrameInfo ] = getattr (task , "__pw_stack__" , inspect .stack ())
430- metadata = _extract_metadata_from_stack (st , is_internal )
431- if metadata :
432- self ._api_zone .set (metadata )
461+ self ._api_zone .set (_extract_stack_trace_information_from_stack (st , is_internal ))
433462 try :
434463 return cb ()
435464 finally :
@@ -444,19 +473,25 @@ def from_nullable_channel(channel: Optional[Channel]) -> Optional[Any]:
444473 return channel ._object if channel else None
445474
446475
447- def _extract_metadata_from_stack (
476+ class StackFrame (TypedDict ):
477+ file : str
478+ line : int
479+ column : int
480+ function : Optional [str ]
481+
482+
483+ class ParsedStackTrace (TypedDict ):
484+ frames : List [StackFrame ]
485+ apiName : Optional [str ]
486+
487+
488+ def _extract_stack_trace_information_from_stack (
448489 st : List [inspect .FrameInfo ], is_internal : bool
449- ) -> Optional [Dict ]:
450- if is_internal :
451- return {
452- "apiName" : "" ,
453- "stack" : [],
454- "internal" : True ,
455- }
490+ ) -> Optional [ParsedStackTrace ]:
456491 playwright_module_path = str (Path (playwright .__file__ ).parents [0 ])
457492 last_internal_api_name = ""
458493 api_name = ""
459- stack : List [Dict ] = []
494+ parsed_frames : List [StackFrame ] = []
460495 for frame in st :
461496 is_playwright_internal = frame .filename .startswith (playwright_module_path )
462497
@@ -466,10 +501,11 @@ def _extract_metadata_from_stack(
466501 method_name += frame [0 ].f_code .co_name
467502
468503 if not is_playwright_internal :
469- stack .append (
504+ parsed_frames .append (
470505 {
471506 "file" : frame .filename ,
472507 "line" : frame .lineno ,
508+ "column" : 0 ,
473509 "function" : method_name ,
474510 }
475511 )
@@ -480,9 +516,8 @@ def _extract_metadata_from_stack(
480516 last_internal_api_name = ""
481517 if not api_name :
482518 api_name = last_internal_api_name
483- if api_name :
484- return {
485- "apiName" : api_name ,
486- "stack" : stack ,
487- }
488- return None
519+
520+ return {
521+ "frames" : parsed_frames ,
522+ "apiName" : "" if is_internal else api_name ,
523+ }
0 commit comments