new tracing and metadata handling

playwright/_impl/_browser_context.py

@@ -48,6 +48,8 @@
 from playwright._impl._frame import Frame
 from playwright._impl._har_router import HarRouter
 from playwright._impl._helper import (
+    HarContentPolicy,
+    HarMode,
     HarRecordingMetadata,
     RouteFromHarNotFoundPolicy,
     RouteHandler,
@@ -61,8 +63,6 @@
     locals_to_params,
     prepare_record_har_options,
     to_impl,
-    HarContentPolicy,
-    HarMode,
 )
 from playwright._impl._network import Request, Response, Route, serialize_headers
 from playwright._impl._page import BindingCall, Page, Worker

playwright/_impl/_connection.py

@@ -14,11 +14,24 @@
 
 import asyncio
 import contextvars
+import datetime
 import inspect
 import sys
 import traceback
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Set,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from greenlet import greenlet
 from pyee import EventEmitter
@@ -33,6 +46,12 @@
     from playwright._impl._playwright import Playwright
 
 
+if sys.version_info >= (3, 8):  # pragma: no cover
+    from typing import TypedDict
+else:  # pragma: no cover
+    from typing_extensions import TypedDict
+
+
 class Channel(AsyncIOEventEmitter):
     def __init__(self, connection: "Connection", guid: str) -> None:
         super().__init__()
@@ -220,10 +239,11 @@ def __init__(
         self._error: Optional[BaseException] = None
         self.is_remote = False
         self._init_task: Optional[asyncio.Task] = None
-        self._api_zone: contextvars.ContextVar[Optional[Dict]] = contextvars.ContextVar(
-            "ApiZone", default=None
-        )
+        self._api_zone: contextvars.ContextVar[
+            Optional[ParsedStackTrace]
+        ] = contextvars.ContextVar("ApiZone", default=None)
         self._local_utils: Optional["LocalUtils"] = local_utils
+        self._stack_collector = MutableSet[List[Dict[str, Any]]]()
 
     @property
     def local_utils(self) -> "LocalUtils":
@@ -271,6 +291,12 @@ def call_on_object_with_known_name(
     ) -> None:
         self._waiting_for_object[guid] = callback
 
+    def start_collecting_call_metadata(self, collector: Any) -> int:
+        return self._stack_collector.add(collector)
+
+    def stop_collecting_call_metadata(self, collector_id: Any) -> None:
+        self._stack_collector.remove_by_id(collector_id)
+
     def _send_message_to_server(
         self, guid: str, method: str, params: Dict
     ) -> ProtocolCallback:
@@ -283,12 +309,30 @@ def _send_message_to_server(
             getattr(task, "__pw_stack_trace__", traceback.extract_stack()),
         )
         self._callbacks[id] = callback
+        stack_trace_information = cast(ParsedStackTrace, self._api_zone.get())
+        for collector in self._stack_collector.all():
+            collector.append({"stack": stack_trace_information["frames"], "id": id})
+        frames = stack_trace_information.get("frames", [])
+        location = (
+            {
+                "file": frames[0]["file"],
+                "line": frames[0]["line"],
+                "column": frames[0]["column"],
+            }
+            if len(frames) > 0
+            else None
+        )
         message = {
             "id": id,
             "guid": guid,
             "method": method,
             "params": self._replace_channels_with_guids(params),
-            "metadata": self._api_zone.get(),
+            "metadata": {
+                "wallTime": int(datetime.datetime.now().timestamp() * 1000),
+                "apiName": stack_trace_information["apiName"],
+                "location": location,
+                "internal": not stack_trace_information["apiName"],
+            },
         }
         self._transport.send(message)
         self._callbacks[id] = callback
@@ -412,9 +456,9 @@ async def wrap_api_call(
             return await cb()
         task = asyncio.current_task(self._loop)
         st: List[inspect.FrameInfo] = getattr(task, "__pw_stack__", inspect.stack())
-        metadata = _extract_metadata_from_stack(st, is_internal)
-        if metadata:
-            self._api_zone.set(metadata)
+        self._api_zone.set(_extract_stack_trace_information_from_stack(
+            st, is_internal
+        ))
         try:
             return await cb()
         finally:
@@ -427,9 +471,9 @@ def wrap_api_call_sync(
             return cb()
         task = asyncio.current_task(self._loop)
         st: List[inspect.FrameInfo] = getattr(task, "__pw_stack__", inspect.stack())
-        metadata = _extract_metadata_from_stack(st, is_internal)
-        if metadata:
-            self._api_zone.set(metadata)
+        self._api_zone.set(_extract_stack_trace_information_from_stack(
+            st, is_internal
+        ))
         try:
             return cb()
         finally:
@@ -444,19 +488,27 @@ def from_nullable_channel(channel: Optional[Channel]) -> Optional[Any]:
     return channel._object if channel else None
 
 
-def _extract_metadata_from_stack(
+class StackFrame(TypedDict):
+    file: str
+    line: int
+    column: int
+    function: Optional[str]
+
+
+class ParsedStackTrace(TypedDict):
+    allFrames: List[StackFrame]
+    frames: List[StackFrame]
+    frameTexts: List[str]
+    apiName: Optional[str]
+
+
+def _extract_stack_trace_information_from_stack(
     st: List[inspect.FrameInfo], is_internal: bool
-) -> Optional[Dict]:
-    if is_internal:
-        return {
-            "apiName": "",
-            "stack": [],
-            "internal": True,
-        }
+) -> Optional[ParsedStackTrace]:
     playwright_module_path = str(Path(playwright.__file__).parents[0])
     last_internal_api_name = ""
     api_name = ""
-    stack: List[Dict] = []
+    parsed_frames: List[Dict] = []
     for frame in st:
         is_playwright_internal = frame.filename.startswith(playwright_module_path)
 
@@ -466,11 +518,15 @@ def _extract_metadata_from_stack(
         method_name += frame[0].f_code.co_name
 
         if not is_playwright_internal:
-            stack.append(
+            parsed_frames.append(
                 {
-                    "file": frame.filename,
-                    "line": frame.lineno,
-                    "function": method_name,
+                    "frame": {
+                        "file": frame.filename,
+                        "line": frame.lineno,
+                        "column": 0,
+                        "function": method_name,
+                    },
+                    "frameText": f"{method_name} at {frame.filename}:{frame.lineno}",
                 }
             )
         if is_playwright_internal:
@@ -480,9 +536,34 @@ def _extract_metadata_from_stack(
             last_internal_api_name = ""
     if not api_name:
         api_name = last_internal_api_name
-    if api_name:
-        return {
-            "apiName": api_name,
-            "stack": stack,
-        }
-    return None
+
+    return {
+        "allFrames": [frame["frame"] for frame in parsed_frames],
+        "frames": [frame["frame"] for frame in parsed_frames],
+        "frameTexts": [frame["frameText"] for frame in parsed_frames],
+        "apiName": "" if is_internal else api_name,
+    }
+
+
+T = TypeVar("T")
+
+
+class MutableSet(Generic[T]):
+    def __init__(self) -> None:
+        self._counter = 0
+        self._set: Set[int] = set()
+        self._set_items: Dict[int, T] = {}
+
+    def add(self, item: T) -> int:
+        self._counter += 1
+        new_id = self._counter
+        self._set.add(new_id)
+        self._set_items[new_id] = item
+        return new_id
+
+    def remove_by_id(self, item_id: int) -> None:
+        self._set.remove(item_id)
+        self._set_items.pop(item_id)
+
+    def all(self) -> List[T]:
+        return list(self._set_items.values())

playwright/_impl/_local_utils.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 
 import base64
-from typing import Dict, List, Optional, cast
+from typing import Dict, Optional, cast
 
-from playwright._impl._api_structures import HeadersArray, NameValue
+from playwright._impl._api_structures import HeadersArray
 from playwright._impl._connection import ChannelOwner
 from playwright._impl._helper import HarLookupResult, locals_to_params
 
@@ -26,8 +26,8 @@ def __init__(
     ) -> None:
         super().__init__(parent, type, guid, initializer)
 
-    async def zip(self, zip_file: str, entries: List[NameValue]) -> None:
-        await self._channel.send("zip", {"zipFile": zip_file, "entries": entries})
+    async def zip(self, params: Dict) -> None:
+        await self._channel.send("zip", params)
 
     async def har_open(self, file: str) -> None:
         params = locals_to_params(locals())

playwright/_impl/_tracing.py

@@ -13,14 +13,18 @@
 # limitations under the License.
 
 import pathlib
-from typing import Dict, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Union, cast
 
 from playwright._impl._artifact import Artifact
 from playwright._impl._connection import ChannelOwner, from_nullable_channel
 from playwright._impl._helper import locals_to_params
 
 
 class Tracing(ChannelOwner):
+    _include_sources = False
+    _metadata_collector: List[Dict[str, Any]] = []
+    _metadata_collector_id: Optional[int] = None
+
     def __init__(
         self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
     ) -> None:
@@ -35,12 +39,21 @@ async def start(
         sources: bool = None,
     ) -> None:
         params = locals_to_params(locals())
+        self._include_sources = bool(sources)
         await self._channel.send("tracingStart", params)
         await self.start_chunk(title)
+        self._metadata_collector = []
+        self._metadata_collector_id = self._connection.start_collecting_call_metadata(
+            self._metadata_collector
+        )
 
     async def start_chunk(self, title: str = None) -> None:
         params = locals_to_params(locals())
         await self._channel.send("tracingStartChunk", params)
+        self._metadata_collector = []
+        self._metadata_collector_id = self._connection.start_collecting_call_metadata(
+            self._metadata_collector
+        )
 
     async def stop_chunk(self, path: Union[pathlib.Path, str] = None) -> None:
         await self._do_stop_chunk(path)
@@ -50,40 +63,60 @@ async def stop(self, path: Union[pathlib.Path, str] = None) -> None:
         await self._channel.send("tracingStop")
 
     async def _do_stop_chunk(self, file_path: Union[pathlib.Path, str] = None) -> None:
+        if self._metadata_collector_id:
+            self._connection.stop_collecting_call_metadata(self._metadata_collector_id)
+        metadata = self._metadata_collector
+        self._metadata_collector = []
+        self._metadata_collector_id = None
+
+        if not file_path:
+            await self._channel.send("tracingStopChunk", {"mode": "discard"})
+            # Not interested in any artifacts
+            return
+
         is_local = not self._connection.is_remote
 
-        mode = "doNotSave"
-        if file_path:
-            if is_local:
-                mode = "compressTraceAndSources"
-            else:
-                mode = "compressTrace"
+        if is_local:
+            entries = await self._channel.send("tracingStopChunk", {"mode": "entries"})
+            await self._connection.local_utils.zip(
+                {
+                    "zipFile": str(file_path),
+                    "entries": entries,
+                    "metadata": metadata,
+                    "mode": "write",
+                    "includeSources": self._include_sources,
+                }
+            )
+            return
 
         result = await self._channel.send_return_as_dict(
             "tracingStopChunk",
             {
-                "mode": mode,
+                "mode": "archive",
             },
         )
-        if not file_path:
-            # Not interested in artifacts.
-            return
 
         artifact = cast(
             Optional[Artifact],
             from_nullable_channel(result.get("artifact")),
         )
 
+        # The artifact may be missing if the browser closed while stopping tracing.
         if not artifact:
-            # The artifact may be missing if the browser closed while stopping tracing.
             return
 
         # Save trace to the final local file.
         await artifact.save_as(file_path)
         await artifact.delete()
 
         # Add local sources to the remote trace if necessary.
-        if result.get("sourceEntries", []):
+        if len(metadata) > 0:
             await self._connection.local_utils.zip(
-                str(file_path), result["sourceEntries"]
+                {
+                    "zipFile": str(file_path),
+                    "entries": result["entries"],
+                    "metadata": metadata,
+                    "mode": "append",
+                    "includeSources": self._include_sources,
+                }
             )