feat: added BrowserType.connect (microsoft#630)

Author: mxschmitt

local-requirements.txt

@@ -11,8 +11,9 @@ pixelmatch==0.2.3
 pre-commit==2.10.1
 pyOpenSSL==20.0.1
 pytest==6.2.2
-pytest-asyncio==0.14.0
+pytest-asyncio==0.15.0
 pytest-cov==2.11.1
+pytest-repeat==0.9.1
 pytest-sugar==0.9.4
 pytest-timeout==1.4.2
 pytest-xdist==2.2.1

playwright/_impl/_browser.py

@@ -50,6 +50,7 @@ def __init__(
         self._is_connected = True
         self._is_closed_or_closing = False
         self._is_remote = False
+        self._is_connected_over_websocket = False
 
         self._contexts: List[BrowserContext] = []
         self._channel.on("close", lambda _: self._on_close())
@@ -59,7 +60,7 @@ def __repr__(self) -> str:
 
     def _on_close(self) -> None:
         self._is_connected = False
-        self.emit(Browser.Events.Disconnected)
+        self.emit(Browser.Events.Disconnected, self)
         self._is_closed_or_closing = True
 
     @property
@@ -153,6 +154,8 @@ async def close(self) -> None:
         except Exception as e:
             if not is_safe_close_error(e):
                 raise e
+        if self._is_connected_over_websocket:
+            await self._connection.stop_async()
 
     @property
     def version(self) -> str:

playwright/_impl/_browser_type.py

@@ -25,6 +25,7 @@
 from playwright._impl._browser_context import BrowserContext
 from playwright._impl._connection import (
     ChannelOwner,
+    Connection,
     from_channel,
     from_nullable_channel,
 )
@@ -35,6 +36,7 @@
     locals_to_params,
     not_installed_error,
 )
+from playwright._impl._transport import WebSocketTransport
 
 
 class BrowserType(ChannelOwner):
@@ -157,6 +159,30 @@ async def connect_over_cdp(
         if default_context:
             browser._contexts.append(default_context)
             default_context._browser = browser
+        return browser
+
+    async def connect(
+        self, ws_endpoint: str, timeout: float = None, slow_mo: float = None
+    ) -> Browser:
+        transport = WebSocketTransport(ws_endpoint, timeout)
+
+        connection = Connection(
+            self._connection._dispatcher_fiber,
+            self._connection._object_factory,
+            transport,
+        )
+        connection._is_sync = self._connection._is_sync
+        connection._loop = self._connection._loop
+        connection._loop.create_task(connection.run())
+        self._connection._child_ws_connections.append(connection)
+        playwright = await connection.wait_for_object_with_known_name("Playwright")
+        pre_launched_browser = playwright._initializer.get("preLaunchedBrowser")
+        assert pre_launched_browser
+        browser = cast(Browser, from_channel(pre_launched_browser))
+        browser._is_remote = True
+        browser._is_connected_over_websocket = True
+
+        transport.once("close", browser._on_close)
 
         return browser
 

playwright/_impl/_connection.py

@@ -44,7 +44,14 @@ async def inner_send(
         if params is None:
             params = {}
         callback = self._connection._send_message_to_server(self._guid, method, params)
-        result = await callback.future
+
+        done, pending = await asyncio.wait(
+            {self._connection._transport.on_error_future, callback.future},
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+        if not callback.future.done():
+            callback.future.cancel()
+        result = next(iter(done)).result()
         # Protocol now has named return values, assume result is one level deeper unless
         # there is explicit ambiguity.
         if not result:
@@ -142,10 +149,13 @@ def __init__(self, connection: "Connection") -> None:
 
 class Connection:
     def __init__(
-        self, dispatcher_fiber: Any, object_factory: Any, driver_executable: Path
+        self,
+        dispatcher_fiber: Any,
+        object_factory: Callable[[ChannelOwner, str, str, Dict], Any],
+        transport: Transport,
     ) -> None:
-        self._dispatcher_fiber: Any = dispatcher_fiber
-        self._transport = Transport(driver_executable)
+        self._dispatcher_fiber = dispatcher_fiber
+        self._transport = transport
         self._transport.on_message = lambda msg: self._dispatch(msg)
         self._waiting_for_object: Dict[str, Any] = {}
         self._last_id = 0
@@ -154,6 +164,7 @@ def __init__(
         self._object_factory = object_factory
         self._is_sync = False
         self._api_name = ""
+        self._child_ws_connections: List["Connection"] = []
 
     async def run_as_sync(self) -> None:
         self._is_sync = True
@@ -165,12 +176,18 @@ async def run(self) -> None:
         await self._transport.run()
 
     def stop_sync(self) -> None:
-        self._transport.stop()
+        self._transport.request_stop()
         self._dispatcher_fiber.switch()
+        self.cleanup()
 
     async def stop_async(self) -> None:
-        self._transport.stop()
+        self._transport.request_stop()
         await self._transport.wait_until_stopped()
+        self.cleanup()
+
+    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) -> Any:
         if guid in self._objects:

playwright/_impl/_page.py

@@ -29,6 +29,7 @@
     ViewportSize,
 )
 from playwright._impl._api_types import Error
+from playwright._impl._artifact import Artifact
 from playwright._impl._connection import (
     ChannelOwner,
     from_channel,
@@ -110,6 +111,7 @@ def __init__(
         self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
     ) -> None:
         super().__init__(parent, type, guid, initializer)
+        self._browser_context: BrowserContext = None  # type: ignore
         self.accessibility = Accessibility(self._channel)
         self.keyboard = Keyboard(self._channel)
         self.mouse = Mouse(self._channel)
@@ -285,7 +287,9 @@ def _on_dialog(self, params: Any) -> None:
     def _on_download(self, params: Any) -> None:
         url = params["url"]
         suggested_filename = params["suggestedFilename"]
-        artifact = from_channel(params["artifact"])
+        artifact = cast(Artifact, from_channel(params["artifact"]))
+        if self._browser_context._browser:
+            artifact._is_remote = self._browser_context._browser._is_remote
         self.emit(
             Page.Events.Download, Download(self, url, suggested_filename, artifact)
         )

playwright/_impl/_transport.py

@@ -17,8 +17,14 @@
 import json
 import os
 import sys
+from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
+
+import websockets
+from pyee import AsyncIOEventEmitter
+
+from playwright._impl._api_types import Error
 
 
 # Sourced from: https://github.com/pytest-dev/pytest/blob/da01ee0a4bb0af780167ecd228ab3ad249511302/src/_pytest/faulthandler.py#L69-L77
@@ -34,15 +40,51 @@ def _get_stderr_fileno() -> Optional[int]:
         return sys.__stderr__.fileno()
 
 
-class Transport:
+class Transport(ABC):
+    def __init__(self) -> None:
+        self.on_message = lambda _: None
+
+    @abstractmethod
+    def request_stop(self) -> None:
+        pass
+
+    def dispose(self) -> None:
+        pass
+
+    @abstractmethod
+    async def wait_until_stopped(self) -> None:
+        pass
+
+    async def run(self) -> None:
+        self._loop = asyncio.get_running_loop()
+        self.on_error_future: asyncio.Future = asyncio.Future()
+
+    @abstractmethod
+    def send(self, message: Dict) -> None:
+        pass
+
+    def serialize_message(self, message: Dict) -> bytes:
+        msg = json.dumps(message)
+        if "DEBUGP" in os.environ:  # pragma: no cover
+            print("\x1b[32mSEND>\x1b[0m", json.dumps(message, indent=2))
+        return msg.encode()
+
+    def deserialize_message(self, data: bytes) -> Any:
+        obj = json.loads(data)
+
+        if "DEBUGP" in os.environ:  # pragma: no cover
+            print("\x1b[33mRECV>\x1b[0m", json.dumps(obj, indent=2))
+        return obj
+
+
+class PipeTransport(Transport):
     def __init__(self, driver_executable: Path) -> None:
         super().__init__()
-        self.on_message = lambda _: None
         self._stopped = False
         self._driver_executable = driver_executable
         self._loop: asyncio.AbstractEventLoop
 
-    def stop(self) -> None:
+    def request_stop(self) -> None:
         self._stopped = True
         self._output.close()
 
@@ -51,7 +93,7 @@ async def wait_until_stopped(self) -> None:
         await self._proc.wait()
 
     async def run(self) -> None:
-        self._loop = asyncio.get_running_loop()
+        await super().run()
         self._stopped_future: asyncio.Future = asyncio.Future()
 
         self._proc = proc = await asyncio.create_subprocess_exec(
@@ -79,21 +121,73 @@ async def run(self) -> None:
                         buffer = buffer + data
                     else:
                         buffer = data
-                obj = json.loads(buffer)
 
-                if "DEBUGP" in os.environ:  # pragma: no cover
-                    print("\x1b[33mRECV>\x1b[0m", json.dumps(obj, indent=2))
+                obj = self.deserialize_message(buffer)
                 self.on_message(obj)
             except asyncio.IncompleteReadError:
                 break
             await asyncio.sleep(0)
         self._stopped_future.set_result(None)
 
     def send(self, message: Dict) -> None:
-        msg = json.dumps(message)
-        if "DEBUGP" in os.environ:  # pragma: no cover
-            print("\x1b[32mSEND>\x1b[0m", json.dumps(message, indent=2))
-        data = msg.encode()
+        data = self.serialize_message(message)
         self._output.write(
             len(data).to_bytes(4, byteorder="little", signed=False) + data
         )
+
+
+class WebSocketTransport(AsyncIOEventEmitter, Transport):
+    def __init__(self, ws_endpoint: str, timeout: float = None) -> None:
+        super().__init__()
+        Transport.__init__(self)
+
+        self._stopped = False
+        self.ws_endpoint = ws_endpoint
+        self.timeout = timeout
+        self._loop: asyncio.AbstractEventLoop
+
+    def request_stop(self) -> None:
+        self._stopped = True
+        self._loop.create_task(self._connection.close())
+
+    def dispose(self) -> None:
+        self.on_error_future.cancel()
+
+    async def wait_until_stopped(self) -> None:
+        await self._connection.wait_closed()
+
+    async def run(self) -> None:
+        await super().run()
+
+        options = {}
+        if self.timeout is not None:
+            options["close_timeout"] = self.timeout / 1000
+            options["ping_timeout"] = self.timeout / 1000
+        self._connection = await websockets.connect(self.ws_endpoint, **options)
+
+        while not self._stopped:
+            try:
+                message = await self._connection.recv()
+                if self._stopped:
+                    self.on_error_future.set_exception(
+                        Error("Playwright connection closed")
+                    )
+                    break
+                obj = self.deserialize_message(message)
+                self.on_message(obj)
+            except websockets.exceptions.ConnectionClosed:
+                if not self._stopped:
+                    self.emit("close")
+                self.on_error_future.set_exception(
+                    Error("Playwright connection closed")
+                )
+                break
+            except Exception as exc:
+                print(f"Received unhandled exception: {exc}")
+                self.on_error_future.set_exception(exc)
+
+    def send(self, message: Dict) -> None:
+        if self._stopped or self._connection.closed:
+            raise Error("Playwright connection closed")
+        data = self.serialize_message(message)
+        self._loop.create_task(self._connection.send(data))

playwright/_impl/_video.py

@@ -28,6 +28,10 @@ def __init__(self, page: "Page") -> None:
         self._dispatcher_fiber = page._dispatcher_fiber
         self._page = page
         self._artifact_future = page._loop.create_future()
+        if page._browser_context and page._browser_context._browser:
+            self._is_remote = page._browser_context._browser._is_remote
+        else:
+            self._is_remote = False
         if page.is_closed():
             self._page_closed()
         else:
@@ -42,9 +46,14 @@ def _page_closed(self) -> None:
 
     def _artifact_ready(self, artifact: Artifact) -> None:
         if not self._artifact_future.done():
+            artifact._is_remote = self._is_remote
             self._artifact_future.set_result(artifact)
 
     async def path(self) -> pathlib.Path:
+        if self._is_remote:
+            raise Error(
+                "Path is not available when using browserType.connect(). Use save_as() to save a local copy."
+            )
         artifact = await self._artifact_future
         if not artifact:
             raise Error("Page did not produce any video frames")

playwright/async_api/_context_manager.py

@@ -18,6 +18,7 @@
 from playwright._impl._connection import Connection
 from playwright._impl._driver import compute_driver_executable
 from playwright._impl._object_factory import create_remote_object
+from playwright._impl._transport import PipeTransport
 from playwright.async_api._generated import Playwright as AsyncPlaywright
 
 
@@ -27,7 +28,7 @@ def __init__(self) -> None:
 
     async def __aenter__(self) -> AsyncPlaywright:
         self._connection = Connection(
-            None, create_remote_object, compute_driver_executable()
+            None, create_remote_object, PipeTransport(compute_driver_executable())
         )
         loop = asyncio.get_running_loop()
         self._connection._loop = loop