feat(websocket): implement Web Sockets (#265)

Author: pavelfeldman

playwright/async_api.py

@@ -62,6 +62,7 @@
 from playwright.network import Request as RequestImpl
 from playwright.network import Response as ResponseImpl
 from playwright.network import Route as RouteImpl
+from playwright.network import WebSocket as WebSocketImpl
 from playwright.page import BindingCall as BindingCallImpl
 from playwright.page import Page as PageImpl
 from playwright.page import Worker as WorkerImpl
@@ -488,6 +489,26 @@ async def continue_(
 mapping.register(RouteImpl, Route)
 
 
+class WebSocket(AsyncBase):
+    def __init__(self, obj: WebSocketImpl):
+        super().__init__(obj)
+
+    @property
+    def url(self) -> str:
+        """WebSocket.url
+
+        Contains the URL of the WebSocket.
+
+        Returns
+        -------
+        str
+        """
+        return mapping.from_maybe_impl(self._impl_obj.url)
+
+
+mapping.register(WebSocketImpl, WebSocket)
+
+
 class Keyboard(AsyncBase):
     def __init__(self, obj: KeyboardImpl):
         super().__init__(obj)

playwright/network.py

@@ -16,6 +16,7 @@
 import json
 import mimetypes
 from pathlib import Path
+from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
 from urllib import parse
 
@@ -257,6 +258,49 @@ def frame(self) -> "Frame":
         return self._request.frame
 
 
+class WebSocket(ChannelOwner):
+
+    Events = SimpleNamespace(
+        Close="close",
+        FrameReceived="framereceived",
+        FrameSent="framesent",
+        Error="socketerror",
+    )
+
+    def __init__(
+        self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
+    ) -> None:
+        super().__init__(parent, type, guid, initializer)
+        self._channel.on(
+            "frameSent",
+            lambda params: self._on_frame_sent(params["opcode"], params["data"]),
+        )
+        self._channel.on(
+            "frameReceived",
+            lambda params: self._on_frame_received(params["opcode"], params["data"]),
+        )
+        self._channel.on(
+            "error", lambda params: self.emit(WebSocket.Events.Error, params["error"])
+        )
+        self._channel.on("close", lambda params: self.emit(WebSocket.Events.Close))
+
+    @property
+    def url(self) -> str:
+        return self._initializer["url"]
+
+    def _on_frame_sent(self, opcode: int, data: str) -> None:
+        if opcode == 2:
+            self.emit(WebSocket.Events.FrameSent, base64.b64decode(data))
+        else:
+            self.emit(WebSocket.Events.FrameSent, data)
+
+    def _on_frame_received(self, opcode: int, data: str) -> None:
+        if opcode == 2:
+            self.emit(WebSocket.Events.FrameReceived, base64.b64decode(data))
+        else:
+            self.emit(WebSocket.Events.FrameReceived, data)
+
+
 def serialize_headers(headers: Dict[str, str]) -> List[Header]:
     return [{"name": name, "value": value} for name, value in headers.items()]
 

playwright/object_factory.py

@@ -26,7 +26,7 @@
 from playwright.element_handle import ElementHandle
 from playwright.frame import Frame
 from playwright.js_handle import JSHandle
-from playwright.network import Request, Response, Route
+from playwright.network import Request, Response, Route, WebSocket
 from playwright.page import BindingCall, Page, Worker
 from playwright.playwright import Playwright
 from playwright.selectors import Selectors
@@ -81,6 +81,8 @@ def create_remote_object(
         return Response(parent, type, guid, initializer)
     if type == "Route":
         return Route(parent, type, guid, initializer)
+    if type == "WebSocket":
+        return WebSocket(parent, type, guid, initializer)
     if type == "Worker":
         return Worker(parent, type, guid, initializer)
     if type == "Selectors":

playwright/page.py

@@ -91,6 +91,7 @@ class Page(ChannelOwner):
         FrameNavigated="framenavigated",
         Load="load",
         Popup="popup",
+        WebSocket="websocket",
         Worker="worker",
     )
     accessibility: Accessibility
@@ -213,6 +214,12 @@ def __init__(
                 params["relativePath"]
             ),
         )
+        self._channel.on(
+            "webSocket",
+            lambda params: self.emit(
+                Page.Events.WebSocket, from_channel(params["webSocket"])
+            ),
+        )
         self._channel.on(
             "worker", lambda params: self._on_worker(from_channel(params["worker"]))
         )

playwright/sync_api.py

@@ -61,6 +61,7 @@
 from playwright.network import Request as RequestImpl
 from playwright.network import Response as ResponseImpl
 from playwright.network import Route as RouteImpl
+from playwright.network import WebSocket as WebSocketImpl
 from playwright.page import BindingCall as BindingCallImpl
 from playwright.page import Page as PageImpl
 from playwright.page import Worker as WorkerImpl
@@ -494,6 +495,26 @@ def continue_(
 mapping.register(RouteImpl, Route)
 
 
+class WebSocket(SyncBase):
+    def __init__(self, obj: WebSocketImpl):
+        super().__init__(obj)
+
+    @property
+    def url(self) -> str:
+        """WebSocket.url
+
+        Contains the URL of the WebSocket.
+
+        Returns
+        -------
+        str
+        """
+        return mapping.from_maybe_impl(self._impl_obj.url)
+
+
+mapping.register(WebSocketImpl, WebSocket)
+
+
 class Keyboard(SyncBase):
     def __init__(self, obj: KeyboardImpl):
         super().__init__(obj)

scripts/expected_api_mismatch.txt

@@ -121,4 +121,3 @@ Parameter not implemented: BrowserContext.waitForEvent(optionsOrPredicate=)
 # 1.6 Todo
 Parameter not implemented: Browser.newPage(proxy=)
 Parameter not implemented: Browser.newContext(proxy=)
-Method not implemented: WebSocket.url

scripts/generate_api.py

@@ -39,7 +39,7 @@
 from playwright.frame import Frame
 from playwright.input import Keyboard, Mouse, Touchscreen
 from playwright.js_handle import JSHandle, Serializable
-from playwright.network import Request, Response, Route
+from playwright.network import Request, Response, Route, WebSocket
 from playwright.page import BindingCall, Page, Worker
 from playwright.playwright import Playwright
 from playwright.selectors import Selectors
@@ -166,7 +166,7 @@ def return_value(value: Any) -> List[str]:
 from playwright.helper import ConsoleMessageLocation, Credentials, MousePosition, Error, FilePayload, SelectOption, RequestFailure, Viewport, DeviceDescriptor, IntSize, FloatRect, Geolocation, ProxyServer, PdfMargins, ResourceTiming, RecordHarOptions
 from playwright.input import Keyboard as KeyboardImpl, Mouse as MouseImpl, Touchscreen as TouchscreenImpl
 from playwright.js_handle import JSHandle as JSHandleImpl
-from playwright.network import Request as RequestImpl, Response as ResponseImpl, Route as RouteImpl
+from playwright.network import Request as RequestImpl, Response as ResponseImpl, Route as RouteImpl, WebSocket as WebSocketImpl
 from playwright.page import BindingCall as BindingCallImpl, Page as PageImpl, Worker as WorkerImpl
 from playwright.playwright import Playwright as PlaywrightImpl
 from playwright.selectors import Selectors as SelectorsImpl
@@ -177,6 +177,7 @@ def return_value(value: Any) -> List[str]:
     Request,
     Response,
     Route,
+    WebSocket,
     Keyboard,
     Mouse,
     Touchscreen,