Skip to content

Commit 3ed5b5c

Browse files
mxschmittmxschmitt
committed
draft
1 parent d9cdfbb commit 3ed5b5c

17 files changed

+1271
-97
lines changed

‎README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ Playwright is a Python library to automate [Chromium](https://www.chromium.org/H
44

55
| | Linux | macOS | Windows |
66
| :--- | :---: | :---: | :---: |
7-
| Chromium <!-- GEN:chromium-version -->129.0.6668.29<!-- GEN:stop --> | ✅ | ✅ | ✅ |
7+
| Chromium <!-- GEN:chromium-version -->130.0.6723.19<!-- GEN:stop --> | ✅ | ✅ | ✅ |
88
| WebKit <!-- GEN:webkit-version -->18.0<!-- GEN:stop --> | ✅ | ✅ | ✅ |
99
| Firefox <!-- GEN:firefox-version -->130.0<!-- GEN:stop --> | ✅ | ✅ | ✅ |
1010

‎playwright/_impl/_browser_context.py

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,14 +62,22 @@
6262
TimeoutSettings,
6363
URLMatch,
6464
URLMatcher,
65+
WebSocketRouteHandlerCallback,
6566
async_readfile,
6667
async_writefile,
6768
locals_to_params,
6869
parse_error,
6970
prepare_record_har_options,
7071
to_impl,
7172
)
72-
from playwright._impl._network import Request, Response, Route, serialize_headers
73+
from playwright._impl._network import (
74+
Request,
75+
Response,
76+
Route,
77+
WebSocketRoute,
78+
WebSocketRouteHandler,
79+
serialize_headers,
80+
)
7381
from playwright._impl._page import BindingCall, Page, Worker
7482
from playwright._impl._str_utils import escape_regex_flags
7583
from playwright._impl._tracing import Tracing
@@ -106,6 +114,7 @@ def __init__(
106114
self._browser._contexts.append(self)
107115
self._pages: List[Page] = []
108116
self._routes: List[RouteHandler] = []
117+
self._web_socket_routes: List[WebSocketRouteHandler] = []
109118
self._bindings: Dict[str, Any] = {}
110119
self._timeout_settings = TimeoutSettings(None)
111120
self._owner_page: Optional[Page] = None
@@ -132,7 +141,14 @@ def __init__(
132141
)
133142
),
134143
)
135-
144+
self._channel.on(
145+
"webSocketRoute",
146+
lambda params: self._loop.create_task(
147+
self._on_web_socket_route(
148+
from_channel(params["webSocketRoute"]),
149+
)
150+
),
151+
)
136152
self._channel.on(
137153
"backgroundPage",
138154
lambda params: self._on_background_page(from_channel(params["page"])),
@@ -248,6 +264,20 @@ async def _on_route(self, route: Route) -> None:
248264
except Exception:
249265
pass
250266

267+
async def _on_web_socket_route(self, web_socket_route: WebSocketRoute) -> None:
268+
route_handler = next(
269+
(
270+
route_handler
271+
for route_handler in self._web_socket_routes
272+
if route_handler.matches(web_socket_route.url)
273+
),
274+
None,
275+
)
276+
if route_handler:
277+
await route_handler.handle(web_socket_route)
278+
else:
279+
web_socket_route.connect_to_server()
280+
251281
def _on_binding(self, binding_call: BindingCall) -> None:
252282
func = self._bindings.get(binding_call._initializer["name"])
253283
if func is None:
@@ -418,6 +448,17 @@ async def _unroute_internal(
418448
return
419449
await asyncio.gather(*map(lambda router: router.stop(behavior), removed)) # type: ignore
420450

451+
async def route_web_socket(
452+
self, url: URLMatch, handler: WebSocketRouteHandlerCallback
453+
) -> None:
454+
self._web_socket_routes.insert(
455+
0,
456+
WebSocketRouteHandler(
457+
URLMatcher(self._options.get("baseURL"), url), handler
458+
),
459+
)
460+
await self._update_web_socket_interception_patterns()
461+
421462
def _dispose_har_routers(self) -> None:
422463
for router in self._har_routers:
423464
router.dispose()
@@ -488,6 +529,14 @@ async def _update_interception_patterns(self) -> None:
488529
"setNetworkInterceptionPatterns", {"patterns": patterns}
489530
)
490531

532+
async def _update_web_socket_interception_patterns(self) -> None:
533+
patterns = WebSocketRouteHandler.prepare_interception_patterns(
534+
self._web_socket_routes
535+
)
536+
await self._channel.send(
537+
"setWebSocketInterceptionPatterns", {"patterns": patterns}
538+
)
539+
491540
def expect_event(
492541
self,
493542
event: str,

‎playwright/_impl/_helper.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,14 +50,15 @@
5050

5151
if TYPE_CHECKING: # pragma: no cover
5252
from playwright._impl._api_structures import HeadersArray
53-
from playwright._impl._network import Request, Response, Route
53+
from playwright._impl._network import Request, Response, Route, WebSocketRoute
5454

5555
URLMatch = Union[str, Pattern[str], Callable[[str], bool]]
5656
URLMatchRequest = Union[str, Pattern[str], Callable[["Request"], bool]]
5757
URLMatchResponse = Union[str, Pattern[str], Callable[["Response"], bool]]
5858
RouteHandlerCallback = Union[
5959
Callable[["Route"], Any], Callable[["Route", "Request"], Any]
6060
]
61+
WebSocketRouteHandlerCallback = Callable[["WebSocketRoute"], Any]
6162

6263
ColorScheme = Literal["dark", "light", "no-preference", "null"]
6364
ForcedColors = Literal["active", "none", "null"]

‎playwright/_impl/_network.py

Lines changed: 217 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import json
1919
import json as json_utils
2020
import mimetypes
21+
import re
2122
from collections import defaultdict
2223
from pathlib import Path
2324
from types import SimpleNamespace
@@ -46,12 +47,19 @@
4647
)
4748
from playwright._impl._connection import (
4849
ChannelOwner,
50+
Connection,
4951
from_channel,
5052
from_nullable_channel,
5153
)
5254
from playwright._impl._errors import Error
5355
from playwright._impl._event_context_manager import EventContextManagerImpl
54-
from playwright._impl._helper import async_readfile, locals_to_params
56+
from playwright._impl._helper import (
57+
URLMatcher,
58+
WebSocketRouteHandlerCallback,
59+
async_readfile,
60+
locals_to_params,
61+
)
62+
from playwright._impl._str_utils import escape_regex_flags
5563
from playwright._impl._waiter import Waiter
5664

5765
if TYPE_CHECKING: # pragma: no cover
@@ -548,6 +556,214 @@ async def _race_with_page_close(self, future: Coroutine) -> None:
548556
await asyncio.gather(fut, return_exceptions=True)
549557

550558

559+
class ServerWebSocketRoute:
560+
def __init__(self, ws: "WebSocketRoute"):
561+
self._ws = ws
562+
563+
def on_message(self, handler: Callable[[Union[str, bytes]], Any]) -> None:
564+
self._ws._on_server_message = handler
565+
566+
def on_close(self, handler: Callable[[Optional[int], Optional[str]], Any]) -> None:
567+
self._ws._on_server_close = handler
568+
569+
def connect_to_server(self) -> None:
570+
raise NotImplementedError(
571+
"connectToServer must be called on the page-side WebSocketRoute"
572+
)
573+
574+
@property
575+
def url(self) -> str:
576+
return self._ws._initializer["url"]
577+
578+
def close(self, code: int = None, reason: str = None) -> None:
579+
try:
580+
asyncio.create_task(
581+
self._ws._channel.send(
582+
"close",
583+
{
584+
"code": code,
585+
"reason": reason,
586+
},
587+
)
588+
)
589+
except:
590+
pass
591+
592+
def send(self, message: Union[str, bytes]) -> None:
593+
if isinstance(message, str):
594+
asyncio.create_task(
595+
self._ws._channel.send(
596+
"sendToServer", {"message": message, "isBase64": False}
597+
)
598+
)
599+
else:
600+
asyncio.create_task(
601+
self._ws._channel.send(
602+
"sendToServer",
603+
{"message": base64.b64encode(message).decode(), "isBase64": True},
604+
)
605+
)
606+
607+
608+
class WebSocketRoute(ChannelOwner):
609+
def __init__(
610+
self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
611+
) -> None:
612+
super().__init__(parent, type, guid, initializer)
613+
self._on_page_message: Optional[Callable[[Union[str, bytes]], Any]] = None
614+
self._on_page_close: Optional[
615+
Callable[[Optional[int], Optional[str]], Any]
616+
] = None
617+
self._on_server_message: Optional[Callable[[Union[str, bytes]], Any]] = None
618+
self._on_server_close: Optional[
619+
Callable[[Optional[int], Optional[str]], Any]
620+
] = None
621+
self._server = ServerWebSocketRoute(self)
622+
self._connected = False
623+
624+
self._channel.on("messageFromPage", self._channel_message_from_page)
625+
self._channel.on("messageFromServer", self._channel_message_from_server)
626+
self._channel.on("closePage", self._channel_close_page)
627+
self._channel.on("closeServer", self._channel_close_server)
628+
629+
def _channel_message_from_page(self, event: Dict) -> None:
630+
if self._on_page_message:
631+
self._on_page_message(
632+
base64.b64decode(event["message"])
633+
if event["isBase64"]
634+
else event["message"]
635+
)
636+
elif self._connected:
637+
try:
638+
asyncio.create_task(self._channel.send("sendToServer", event))
639+
except:
640+
pass
641+
642+
def _channel_message_from_server(self, event: Dict) -> None:
643+
if self._on_server_message:
644+
self._on_server_message(
645+
base64.b64decode(event["message"])
646+
if event["isBase64"]
647+
else event["message"]
648+
)
649+
else:
650+
try:
651+
asyncio.create_task(self._channel.send("sendToPage", event))
652+
except:
653+
pass
654+
655+
def _channel_close_page(self, event: Dict) -> None:
656+
if self._on_page_close:
657+
self._on_page_close(event["code"], event["reason"])
658+
else:
659+
try:
660+
asyncio.create_task(self._channel.send("closeServer", event))
661+
except:
662+
pass
663+
664+
def _channel_close_server(self, event: Dict) -> None:
665+
if self._on_server_close:
666+
self._on_server_close(event["code"], event["reason"])
667+
else:
668+
try:
669+
asyncio.create_task(self._channel.send("closePage", event))
670+
except:
671+
pass
672+
673+
@property
674+
def url(self) -> str:
675+
return self._initializer["url"]
676+
677+
async def close(self, code: int = None, reason: str = None) -> None:
678+
try:
679+
await self._channel.send(
680+
"closePage", {"code": code, "reason": reason, "wasClean": True}
681+
)
682+
except:
683+
pass
684+
685+
def connect_to_server(self) -> "WebSocketRoute":
686+
if self._connected:
687+
raise Error("Already connected to the server")
688+
self._connected = True
689+
asyncio.create_task(self._channel.send("connect"))
690+
return cast("WebSocketRoute", self._server)
691+
692+
def send(self, message: Union[str, bytes]) -> None:
693+
if isinstance(message, str):
694+
try:
695+
asyncio.create_task(
696+
self._channel.send(
697+
"sendToPage", {"message": message, "isBase64": False}
698+
)
699+
)
700+
except:
701+
pass
702+
else:
703+
try:
704+
asyncio.create_task(
705+
self._channel.send(
706+
"sendToPage",
707+
{
708+
"message": base64.b64encode(message).decode(),
709+
"isBase64": True,
710+
},
711+
)
712+
)
713+
except:
714+
pass
715+
716+
def on_message(self, handler: Callable[[Union[str, bytes]], Any]) -> None:
717+
self._on_page_message = handler
718+
719+
def on_close(self, handler: Callable[[Optional[int], Optional[str]], Any]) -> None:
720+
self._on_page_close = handler
721+
722+
async def _after_handle(self) -> None:
723+
if self._connected:
724+
return
725+
# Ensure that websocket is "open" and can send messages without an actual server connection.
726+
await self._channel.send("ensureOpened")
727+
728+
729+
class WebSocketRouteHandler:
730+
def __init__(self, matcher: URLMatcher, handler: WebSocketRouteHandlerCallback):
731+
self.matcher = matcher
732+
self.handler = handler
733+
734+
@staticmethod
735+
def prepare_interception_patterns(
736+
handlers: List["WebSocketRouteHandler"],
737+
) -> List[dict]:
738+
patterns = []
739+
all_urls = False
740+
for handler in handlers:
741+
if isinstance(handler.matcher.match, str):
742+
patterns.append({"glob": handler.matcher.match})
743+
elif isinstance(handler.matcher._regex_obj, re.Pattern):
744+
patterns.append(
745+
{
746+
"regexSource": handler.matcher._regex_obj.pattern,
747+
"regexFlags": escape_regex_flags(handler.matcher._regex_obj),
748+
}
749+
)
750+
else:
751+
all_urls = True
752+
753+
if all_urls:
754+
return [{"glob": "**/*"}]
755+
return patterns
756+
757+
def matches(self, ws_url: str) -> bool:
758+
return self.matcher.matches(ws_url)
759+
760+
async def handle(self, websocket_route: "WebSocketRoute") -> None:
761+
maybe_future = self.handler(websocket_route)
762+
if maybe_future:
763+
breakpoint()
764+
await websocket_route._after_handle()
765+
766+
551767
class Response(ChannelOwner):
552768
def __init__(
553769
self, parent: ChannelOwner, type: str, guid: str, initializer: Dict

‎playwright/_impl/_object_factory.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,13 @@
2626
from playwright._impl._frame import Frame
2727
from playwright._impl._js_handle import JSHandle
2828
from playwright._impl._local_utils import LocalUtils
29-
from playwright._impl._network import Request, Response, Route, WebSocket
29+
from playwright._impl._network import (
30+
Request,
31+
Response,
32+
Route,
33+
WebSocket,
34+
WebSocketRoute,
35+
)
3036
from playwright._impl._page import BindingCall, Page, Worker
3137
from playwright._impl._playwright import Playwright
3238
from playwright._impl._selectors import SelectorsOwner
@@ -88,6 +94,8 @@ def create_remote_object(
8894
return Tracing(parent, type, guid, initializer)
8995
if type == "WebSocket":
9096
return WebSocket(parent, type, guid, initializer)
97+
if type == "WebSocketRoute":
98+
return WebSocketRoute(parent, type, guid, initializer)
9199
if type == "Worker":
92100
return Worker(parent, type, guid, initializer)
93101
if type == "WritableStream":

0 commit comments

Comments
 (0)