feat: WsPipeTransport

Author: ProgramRipper

playwright/_impl/_browser_type.py

@@ -44,9 +44,9 @@
     async_readfile,
     locals_to_params,
 )
-from playwright._impl._json_pipe import JsonPipeTransport
 from playwright._impl._network import serialize_headers, to_client_certificates_protocol
 from playwright._impl._waiter import throw_on_timeout
+from playwright._impl._ws_pipe import WsPipeTransport
 
 if TYPE_CHECKING:
     from playwright._impl._playwright import Playwright
@@ -224,21 +224,14 @@ async def connect(
             slowMo = 0
 
         headers = {**(headers if headers else {}), "x-playwright-browser": self.name}
-        local_utils = self._connection.local_utils
-        pipe_channel = (
-            await local_utils._channel.send_return_as_dict(
-                "connect",
-                None,
-                {
-                    "wsEndpoint": wsEndpoint,
-                    "headers": headers,
-                    "slowMo": slowMo,
-                    "timeout": timeout if timeout is not None else 0,
-                    "exposeNetwork": exposeNetwork,
-                },
-            )
-        )["pipe"]
-        transport = JsonPipeTransport(self._connection._loop, pipe_channel)
+        transport = WsPipeTransport(
+            self._connection._loop,
+            ws_endpoint=wsEndpoint,
+            headers=headers,
+            slow_mo=slowMo,
+            timeout=timeout if timeout is not None else 0,
+            expose_network=exposeNetwork,
+        )
 
         connection = Connection(
             self._connection._dispatcher_fiber,

playwright/_impl/_connection.py

@@ -367,7 +367,6 @@ def _send_message_to_server(
             getattr(task, "__pw_stack_trace__", traceback.extract_stack(limit=10)),
         )
         callback.no_reply = no_reply
-        self._callbacks[id] = callback
         stack_trace_information = cast(ParsedStackTrace, self._api_zone.get())
         frames = stack_trace_information.get("frames", [])
         location = (
@@ -399,8 +398,8 @@ def _send_message_to_server(
         if self._tracing_count > 0 and frames and object._guid != "localUtils":
             self.local_utils.add_stack_to_tracing_no_reply(id, frames)
 
-        self._transport.send(message)
         self._callbacks[id] = callback
+        self._transport.send(message)
 
         return callback
 

playwright/_impl/_fake_pipe.py

@@ -0,0 +1,71 @@
+import asyncio
+from pathlib import Path
+from typing import Dict
+
+from playwright._impl._driver import compute_driver_executable, get_driver_env
+from playwright._impl._transport import Transport
+
+
+class FakePipeTransport(Transport):
+    def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
+        super().__init__(loop)
+
+    def request_stop(self) -> None:
+        self._stopped_future.set_result(None)
+
+    async def wait_until_stopped(self) -> None:
+        await self._stopped_future
+
+    async def connect(self) -> None:
+        self._stopped_future = asyncio.Future()
+
+    async def run(self) -> None:
+        await self._stopped_future
+
+    def send(self, message: Dict) -> None:
+        if message["method"] != "initialize":
+            try:
+                for path in compute_driver_executable():
+                    Path(path).stat()
+            except FileNotFoundError as e:
+                return self.on_error_future.set_exception(e)
+
+            return self.on_error_future.set_exception(FileNotFoundError)
+
+        for type_, initializer, guid in [
+            (
+                "BrowserType",
+                {"executablePath": "", "name": name},
+                f"browser-type@{name}",
+            )
+            for name in ["chromium", "firefox", "webkit"]
+        ] + [
+            ("LocalUtils", {"deviceDescriptors": []}, "localUtils"),
+            (
+                "Playwright",
+                {
+                    name: {"guid": f"browser-type@{name}"}
+                    for name in ["chromium", "firefox", "webkit"]
+                }
+                | {
+                    "utils": {"guid": "localUtils"},
+                },
+                "Playwright",
+            ),
+        ]:
+            self.on_message(
+                {
+                    "guid": "",
+                    "method": "__create__",
+                    "params": {
+                        "type": type_,
+                        "initializer": initializer,
+                        "guid": guid,
+                    },
+                }
+            )
+
+        self.on_message(
+            {"id": message["id"], "result": {"playwright": {"guid": "Playwright"}}}
+        )
+

playwright/_impl/_ws_pipe.py

@@ -0,0 +1,75 @@
+import asyncio
+from typing import Dict
+
+from pyee.asyncio import AsyncIOEventEmitter
+from websockets import ConnectionClosed, connect
+
+from playwright._impl._transport import Transport
+
+
+class WsPipeTransport(AsyncIOEventEmitter, Transport):
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        *,
+        ws_endpoint: str,
+        timeout: float = None,
+        slow_mo: float = None,
+        headers: Dict[str, str] = None,
+        expose_network: str = None,
+    ) -> None:
+        super().__init__(loop)
+        Transport.__init__(self, loop)
+        self._stop_event = asyncio.Event()
+
+        self._ws_endpoint = ws_endpoint
+        self._timeout = timeout
+        self._slow_mo = slow_mo
+        self._headers = headers
+        self._expose_network = expose_network
+
+    def request_stop(self) -> None:
+        self._stop_event.set()
+
+    def dispose(self) -> None:
+        self.on_error_future.cancel()
+
+    async def wait_until_stopped(self) -> None:
+        await self._ws.wait_closed()
+
+    async def connect(self) -> None:
+        try:
+            self._ws = await connect(
+                self._ws_endpoint, additional_headers=self._headers, max_size=None
+            )
+        except Exception as e:
+            self.on_error_future.set_exception(e)
+            raise e
+
+    async def run(self) -> None:
+        wait = asyncio.create_task(self._stop_event.wait())
+
+        while not self._stop_event.is_set():
+            recv = asyncio.create_task(self._ws.recv(False))
+            await asyncio.wait([recv, wait], return_when=asyncio.FIRST_COMPLETED)
+
+            if self._stop_event.is_set():
+                recv.cancel()
+                break
+
+            try:
+                self.on_message(self.deserialize_message(recv.result()))
+            except ConnectionClosed:
+                if not self._stop_event.is_set():
+                    self.on_error_future.set_exception(
+                        Exception("Connection closed while reading from the driver")
+                    )
+                break
+
+        wait.cancel()
+        self.emit("close", "")
+        await self._ws.close()
+
+    def send(self, message: Dict) -> None:
+        self._ws.protocol.send_binary(self.serialize_message(message))
+        self._ws.send_data()

playwright/async_api/_context_manager.py

@@ -16,8 +16,8 @@
 from typing import Any
 
 from playwright._impl._connection import Connection
+from playwright._impl._fake_pipe import FakePipeTransport
 from playwright._impl._object_factory import create_remote_object
-from playwright._impl._transport import PipeTransport
 from playwright.async_api._generated import Playwright as AsyncPlaywright
 
 
@@ -31,7 +31,7 @@ async def __aenter__(self) -> AsyncPlaywright:
         self._connection = Connection(
             None,
             create_remote_object,
-            PipeTransport(loop),
+            FakePipeTransport(loop),
             loop,
         )
         loop.create_task(self._connection.run())

pyproject.toml

@@ -17,7 +17,8 @@ requires-python = ">=3.9"
 # - uv pip compile pyproject.toml -o requirements.txt
 dependencies = [
     "pyee>=13,<14",
-    "greenlet>=3.1.1,<4.0.0"
+    "greenlet>=3.1.1,<4.0.0",
+    "websockets>=14,<16"
 ]
 classifiers = [
     "Topic :: Software Development :: Testing",

requirements.txt

@@ -6,3 +6,5 @@ pyee==13.0.0
     # via playwright (pyproject.toml)
 typing-extensions==4.14.1
     # via pyee
+websockets==15.0.1
+    # via playwright (pyproject.toml)

setup.py

@@ -72,6 +72,12 @@
         "platform": "win32",
         "zip_name": "win32_arm64",
     },
+    {
+        "wheel": "any.whl",
+        "machine": platform.machine().lower(),
+        "platform": sys.platform,
+        "zip_name": "",
+    },
 ]
 
 if len(sys.argv) == 2 and sys.argv[1] == "--list-wheels":
@@ -120,9 +126,6 @@ def download_driver(zip_name: str) -> None:
 class PlaywrightBDistWheelCommand(BDistWheelCommand):
     def run(self) -> None:
         super().run()
-        os.makedirs("driver", exist_ok=True)
-        os.makedirs("playwright/driver", exist_ok=True)
-        self._download_and_extract_local_driver()
 
         wheel = None
         if os.getenv("PLAYWRIGHT_TARGET_WHEEL", None):
@@ -142,7 +145,11 @@ def run(self) -> None:
                 )
             )[0]
         assert wheel
-        self._build_wheel(wheel)
+        if wheel["zip_name"]:
+            os.makedirs("driver", exist_ok=True)
+            os.makedirs("playwright/driver", exist_ok=True)
+            self._download_and_extract_local_driver()
+            self._build_wheel(wheel)
 
     def _build_wheel(
         self,