shanxia2.0/backend/venv/lib/python3.9/site-packages/playwright/_impl/_tracing.py

# Copyright (c) Microsoft Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pathlib
from typing import Any, Dict, Literal, Optional, Pattern, Union, cast

from playwright._impl._api_structures import TracingGroupLocation
from playwright._impl._artifact import Artifact
from playwright._impl._connection import (
    ChannelOwner,
    from_channel,
    from_nullable_channel,
)
from playwright._impl._disposable import DisposableStub
from playwright._impl._helper import Error, locals_to_params


class Tracing(ChannelOwner):
    def __init__(
        self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
    ) -> None:
        super().__init__(parent, type, guid, initializer)
        self._include_sources: bool = False
        self._is_live: bool = False
        self._stacks_id: Optional[str] = None
        self._is_tracing: bool = False
        self._traces_dir: Optional[str] = None
        self._har_id: Optional[str] = None
        self._har_recorders: Dict[str, Dict[str, str]] = {}

    async def start(
        self,
        name: str = None,
        title: str = None,
        snapshots: bool = None,
        screenshots: bool = None,
        sources: bool = None,
        live: bool = None,
    ) -> None:
        params = locals_to_params(locals())
        self._include_sources = bool(sources)
        self._is_live = bool(live)

        await self._channel.send(
            "tracingStart",
            None,
            {
                "name": name,
                "snapshots": snapshots,
                "screenshots": screenshots,
                "live": live,
            },
        )
        trace_name = await self._channel.send(
            "tracingStartChunk", None, {"title": title, "name": name}
        )
        await self._start_collecting_stacks(trace_name)

    async def start_chunk(self, title: str = None, name: str = None) -> None:
        params = locals_to_params(locals())
        trace_name = await self._channel.send("tracingStartChunk", None, params)
        await self._start_collecting_stacks(trace_name)

    async def _start_collecting_stacks(self, trace_name: str) -> None:
        if not self._is_tracing:
            self._is_tracing = True
            self._connection.set_is_tracing(True)
        self._stacks_id = await self._connection.local_utils.tracing_started(
            self._traces_dir, trace_name, self._is_live
        )

    async def stop_chunk(self, path: Union[pathlib.Path, str] = None) -> None:
        await self._do_stop_chunk(path)

    async def stop(self, path: Union[pathlib.Path, str] = None) -> None:
        await self._do_stop_chunk(path)
        await self._channel.send(
            "tracingStop",
            None,
        )

    async def _do_stop_chunk(self, file_path: Union[pathlib.Path, str] = None) -> None:
        self._reset_stack_counter()

        if not file_path:
            # Not interested in any artifacts
            await self._channel.send("tracingStopChunk", None, {"mode": "discard"})
            if self._stacks_id:
                await self._connection.local_utils.trace_discarded(self._stacks_id)
            return

        is_local = not self._connection.is_remote

        if is_local:
            result = await self._channel.send_return_as_dict(
                "tracingStopChunk", None, {"mode": "entries"}
            )
            await self._connection.local_utils.zip(
                {
                    "zipFile": str(file_path),
                    "entries": result["entries"],
                    "stacksId": self._stacks_id,
                    "mode": "write",
                    "includeSources": self._include_sources,
                }
            )
            return

        result = await self._channel.send_return_as_dict(
            "tracingStopChunk",
            None,
            {
                "mode": "archive",
            },
        )

        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:
            if self._stacks_id:
                await self._connection.local_utils.trace_discarded(self._stacks_id)
            return

        # Save trace to the final local file.
        await artifact.save_as(file_path)
        await artifact.delete()

        await self._connection.local_utils.zip(
            {
                "zipFile": str(file_path),
                "entries": [],
                "stacksId": self._stacks_id,
                "mode": "append",
                "includeSources": self._include_sources,
            }
        )

    def _reset_stack_counter(self) -> None:
        if self._is_tracing:
            self._is_tracing = False
            self._connection.set_is_tracing(False)

    async def group(
        self, name: str, location: TracingGroupLocation = None
    ) -> DisposableStub:
        await self._channel.send("tracingGroup", None, locals_to_params(locals()))
        return DisposableStub(lambda: self.group_end(), self)

    async def group_end(self) -> None:
        await self._channel.send(
            "tracingGroupEnd",
            None,
        )

    async def start_har(
        self,
        path: Union[pathlib.Path, str],
        content: Literal["attach", "embed", "omit"] = None,
        mode: Literal["full", "minimal"] = None,
        urlFilter: Union[str, Pattern[str]] = None,
    ) -> DisposableStub:
        if self._har_id:
            raise Error("HAR recording has already been started")
        is_zip = str(path).endswith(".zip")
        default_content: Literal["attach", "embed", "omit"] = (
            "attach" if is_zip else "embed"
        )
        self._har_id = await self._record_into_har(
            har=path,
            page=None,
            url=urlFilter,
            update_content=content or default_content,
            update_mode=mode or "full",
        )
        return DisposableStub(lambda: self.stop_har(), self)

    async def _record_into_har(
        self,
        har: Union[pathlib.Path, str],
        page: Optional[ChannelOwner],
        url: Union[str, Pattern[str]] = None,
        update_content: Literal["attach", "embed", "omit"] = None,
        update_mode: Literal["full", "minimal"] = None,
        resourcesDir: Optional[str] = None,
    ) -> str:
        is_zip = str(har).endswith(".zip")
        url_glob: Optional[str] = None
        url_regex_source: Optional[str] = None
        url_regex_flags: Optional[str] = None
        if isinstance(url, str):
            url_glob = url
        elif url is not None:
            url_regex_source = url.pattern
            url_regex_flags = "".join(
                flag
                for flag, mask in (("i", 2), ("m", 8), ("s", 16))
                if url.flags & mask
            )
        options: Dict[str, object] = {
            "content": update_content or "attach",
            "mode": update_mode or "minimal",
            "harPath": None if is_zip else str(har),
        }
        if url_glob is not None:
            options["urlGlob"] = url_glob
        if url_regex_source is not None:
            options["urlRegexSource"] = url_regex_source
        if url_regex_flags is not None:
            options["urlRegexFlags"] = url_regex_flags
        if resourcesDir is not None:
            options["resourcesDir"] = resourcesDir
        params: Dict[str, Any] = {"options": options}
        if page is not None:
            params["page"] = page._channel
        result = await self._channel.send_return_as_dict("harStart", None, params)
        har_id = result["harId"]
        self._har_recorders[har_id] = {"path": str(har)}
        return har_id

    async def _export_all_hars(self) -> None:
        for har_id in list(self._har_recorders.keys()):
            await self._export_har(har_id)
        self._har_id = None

    async def stop_har(self) -> None:
        har_id = self._har_id
        if not har_id:
            return
        self._har_id = None
        await self._export_har(har_id)

    async def _export_har(self, har_id: str) -> None:
        params = self._har_recorders.pop(har_id, None)
        if not params:
            return
        is_local = not self._connection.is_remote
        is_zip = params["path"].endswith(".zip")

        if is_local:
            result = await self._channel.send_return_as_dict(
                "harExport", None, {"harId": har_id, "mode": "entries"}
            )
            if not is_zip:
                # Server wrote HAR and resources to the user's chosen paths.
                return
            await self._connection.local_utils.zip(
                {
                    "zipFile": params["path"],
                    "entries": result["entries"],
                    "mode": "write",
                    "includeSources": False,
                }
            )
            return

        result = await self._channel.send_return_as_dict(
            "harExport", None, {"harId": har_id, "mode": "archive"}
        )
        artifact = cast(Artifact, from_channel(result["artifact"]))
        if is_zip:
            await artifact.save_as(params["path"])
            await artifact.delete()
            return
        # Uncompressed har is not supported in thin clients
        await artifact.save_as(params["path"] + ".tmp")
        await artifact.delete()