chore(protocol): implement client side hello

mxschmitt · mxschmitt · commit 2de45fba5b6a · 2021-08-19T19:50:09.000+02:00
diff --git a/playwright/_impl/_browser_type.py b/playwright/_impl/_browser_type.py
@@ -15,7 +15,7 @@
 import asyncio
 import pathlib
 from pathlib import Path
-from typing import Dict, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Dict, List, Optional, Union, cast
 
 from playwright._impl._api_structures import (
     Geolocation,
@@ -42,6 +42,9 @@
 from playwright._impl._transport import WebSocketTransport
 from playwright._impl._wait_helper import throw_on_timeout
 
+if TYPE_CHECKING:
+    from playwright._impl._playwright import Playwright
+
 
 class BrowserType(ChannelOwner):
     def __init__(
@@ -183,8 +186,13 @@ async def connect(
         if timeout is None:
             timeout = 30000
 
+        on_ready = asyncio.Event()
         transport = WebSocketTransport(
-            self._connection._loop, ws_endpoint, headers, slow_mo
+            self._connection._loop,
+            ws_endpoint,
+            headers,
+            slow_mo,
+            lambda: on_ready.set(),
         )
         connection = Connection(
             self._connection._dispatcher_fiber,
@@ -194,9 +202,12 @@ async def connect(
         connection._is_sync = self._connection._is_sync
         connection._loop = self._connection._loop
         connection._loop.create_task(connection.run())
-        future = connection._loop.create_task(
-            connection.wait_for_object_with_known_name("Playwright")
-        )
+
+        async def initialize() -> "Playwright":
+            await on_ready.wait()
+            return await connection.initialize_playwright()
+
+        future = asyncio.create_task(initialize())
         timeout_future = throw_on_timeout(timeout, Error("Connection timed out"))
         done, pending = await asyncio.wait(
             {transport.on_error_future, future, timeout_future},
@@ -206,7 +217,7 @@ async def connect(
             future.cancel()
         if not timeout_future.done():
             timeout_future.cancel()
-        playwright = next(iter(done)).result()
+        playwright: "Playwright" = next(iter(done)).result()
         self._connection._child_ws_connections.append(connection)
         pre_launched_browser = playwright._initializer.get("preLaunchedBrowser")
         assert pre_launched_browser
diff --git a/playwright/_impl/_connection.py b/playwright/_impl/_connection.py
@@ -16,14 +16,17 @@
 import sys
 import traceback
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from greenlet import greenlet
 from pyee import AsyncIOEventEmitter
 
 from playwright._impl._helper import ParsedMessagePayload, parse_error
 from playwright._impl._transport import Transport
 
+if TYPE_CHECKING:
+    from playwright._impl._playwright import Playwright
+
 
 class Channel(AsyncIOEventEmitter):
     def __init__(self, connection: "Connection", guid: str) -> None:
@@ -119,7 +122,7 @@ def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
 
 class RootChannelOwner(ChannelOwner):
     def __init__(self, connection: "Connection") -> None:
-        super().__init__(connection, "", "", {})
+        super().__init__(connection, "Root", "", {})
 
 
 class Connection:
@@ -164,16 +167,15 @@ def cleanup(self) -> None:
         for ws_connection in self._child_ws_connections:
             ws_connection._transport.dispose()
 
-    async def wait_for_object_with_known_name(self, guid: str) -> ChannelOwner:
-        if guid in self._objects:
-            return self._objects[guid]
-        callback: asyncio.Future[ChannelOwner] = self._loop.create_future()
-
-        def callback_wrapper(result: ChannelOwner) -> None:
-            callback.set_result(result)
-
-        self._waiting_for_object[guid] = callback_wrapper
-        return await callback
+    async def initialize_playwright(self) -> "Playwright":
+        result = await self._send_message_to_server(
+            "",
+            "initialize",
+            {
+                "language": "python",
+            },
+        ).future
+        return from_channel(result["playwright"])
 
     def call_on_object_with_known_name(
         self, guid: str, callback: Callable[[ChannelOwner], None]
diff --git a/playwright/_impl/_transport.py b/playwright/_impl/_transport.py
@@ -44,10 +44,11 @@ def _get_stderr_fileno() -> Optional[int]:
 
 
 class Transport(ABC):
-    def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
+    def __init__(self, loop: asyncio.AbstractEventLoop, on_ready: Callable) -> None:
         self._loop = loop
         self.on_message: Callable[[ParsedMessagePayload], None] = lambda _: None
         self.on_error_future: asyncio.Future = loop.create_future()
+        self._on_ready = on_ready
 
     @abstractmethod
     def request_stop(self) -> None:
@@ -84,9 +85,12 @@ def deserialize_message(self, data: Union[str, bytes]) -> ParsedMessagePayload:
 
 class PipeTransport(Transport):
     def __init__(
-        self, loop: asyncio.AbstractEventLoop, driver_executable: Path
+        self,
+        loop: asyncio.AbstractEventLoop,
+        driver_executable: Path,
+        on_ready: Callable,
     ) -> None:
-        super().__init__(loop)
+        super().__init__(loop, on_ready)
         self._stopped = False
         self._driver_executable = driver_executable
 
@@ -128,6 +132,7 @@ async def run(self) -> None:
         assert proc.stdout
         assert proc.stdin
         self._output = proc.stdin
+        self._on_ready()
 
         while not self._stopped:
             try:
@@ -162,16 +167,18 @@ def __init__(
         self,
         loop: asyncio.AbstractEventLoop,
         ws_endpoint: str,
-        headers: Dict[str, str] = None,
-        slow_mo: float = None,
+        headers: Optional[Dict[str, str]],
+        slow_mo: Optional[float],
+        on_ready: Callable,
     ) -> None:
-        super().__init__(loop)
-        Transport.__init__(self, loop)
+        AsyncIOEventEmitter.__init__(self)
+        Transport.__init__(self, loop, on_ready)
 
         self._stopped = False
         self.ws_endpoint = ws_endpoint
         self.headers = headers
         self.slow_mo = slow_mo
+        self._on_ready = on_ready
 
     def request_stop(self) -> None:
         self._stopped = True
@@ -192,7 +199,7 @@ async def run(self) -> None:
         except Exception as exc:
             self.on_error_future.set_exception(Error(f"websocket.connect: {str(exc)}"))
             return
-
+        self._on_ready()
         while not self._stopped:
             try:
                 message = await self._connection.recv()
diff --git a/playwright/async_api/_context_manager.py b/playwright/async_api/_context_manager.py
@@ -28,15 +28,22 @@ def __init__(self) -> None:
 
     async def __aenter__(self) -> AsyncPlaywright:
         loop = asyncio.get_running_loop()
+        on_ready = asyncio.Event()
         self._connection = Connection(
             None,
             create_remote_object,
-            PipeTransport(loop, compute_driver_executable()),
+            PipeTransport(
+                loop,
+                compute_driver_executable(),
+                lambda: on_ready.set(),
+            ),
         )
         self._connection._loop = loop
         loop.create_task(self._connection.run())
-        playwright_future = loop.create_task(
-            self._connection.wait_for_object_with_known_name("Playwright")
+        await on_ready.wait()
+
+        playwright_future = asyncio.create_task(
+            self._connection.initialize_playwright()
         )
         done, pending = await asyncio.wait(
             {self._connection._transport.on_error_future, playwright_future},
diff --git a/playwright/sync_api/_context_manager.py b/playwright/sync_api/_context_manager.py
@@ -52,10 +52,14 @@ def greenlet_main() -> None:
                 loop.close()
 
         dispatcher_fiber = greenlet(greenlet_main)
+
+        def on_ready() -> None:
+            asyncio.ensure_future(self._connection.initialize_playwright())
+
         self._connection = Connection(
             dispatcher_fiber,
             create_remote_object,
-            PipeTransport(loop, compute_driver_executable()),
+            PipeTransport(loop, compute_driver_executable(), on_ready),
         )
 
         g_self = greenlet.getcurrent()
diff --git a/setup.py b/setup.py
@@ -28,7 +28,7 @@
     InWheel = None
 from wheel.bdist_wheel import bdist_wheel as BDistWheelCommand
 
-driver_version = "1.14.0-1628783206000"
+driver_version = "1.15.0-next-1629387074000"
 
 
 def extractall(zip: zipfile.ZipFile, path: str) -> None:
