chore: align selectors handling with upstream (microsoft#1176)

Author: mxschmitt

playwright/_impl/_browser_type.py

@@ -206,6 +206,7 @@ async def connect(
         if not timeout_future.done():
             timeout_future.cancel()
         playwright: "Playwright" = next(iter(done)).result()
+        playwright._set_selectors(self._playwright.selectors)
         self._connection._child_ws_connections.append(connection)
         pre_launched_browser = playwright._initializer.get("preLaunchedBrowser")
         assert pre_launched_browser
@@ -219,6 +220,7 @@ def handle_transport_close() -> None:
                     page._on_close()
                 context._on_close()
             browser._on_close()
+            connection.cleanup()
 
         transport.once("close", handle_transport_close)
 

playwright/_impl/_connection.py

@@ -19,7 +19,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from greenlet import greenlet
-from pyee import AsyncIOEventEmitter
+from pyee import AsyncIOEventEmitter, EventEmitter
 
 from playwright._impl._helper import ParsedMessagePayload, parse_error
 from playwright._impl._transport import Transport
@@ -141,14 +141,15 @@ async def initialize(self) -> "Playwright":
         )
 
 
-class Connection:
+class Connection(EventEmitter):
     def __init__(
         self,
         dispatcher_fiber: Any,
         object_factory: Callable[[ChannelOwner, str, str, Dict], ChannelOwner],
         transport: Transport,
         loop: asyncio.AbstractEventLoop,
     ) -> None:
+        super().__init__()
         self._dispatcher_fiber = dispatcher_fiber
         self._transport = transport
         self._transport.on_message = lambda msg: self.dispatch(msg)
@@ -195,6 +196,7 @@ async def stop_async(self) -> None:
     def cleanup(self) -> None:
         for ws_connection in self._child_ws_connections:
             ws_connection._transport.dispose()
+        self.emit("close")
 
     def call_on_object_with_known_name(
         self, guid: str, callback: Callable[[ChannelOwner], None]

playwright/_impl/_object_factory.py

@@ -30,7 +30,7 @@
 from playwright._impl._network import Request, Response, Route, WebSocket
 from playwright._impl._page import BindingCall, Page, Worker
 from playwright._impl._playwright import Playwright
-from playwright._impl._selectors import Selectors
+from playwright._impl._selectors import SelectorsOwner
 from playwright._impl._stream import Stream
 from playwright._impl._tracing import Tracing
 
@@ -90,5 +90,5 @@ def create_remote_object(
     if type == "Worker":
         return Worker(parent, type, guid, initializer)
     if type == "Selectors":
-        return Selectors(parent, type, guid, initializer)
+        return SelectorsOwner(parent, type, guid, initializer)
     return DummyObject(parent, type, guid, initializer)

playwright/_impl/_playwright.py

@@ -18,7 +18,7 @@
 from playwright._impl._connection import ChannelOwner, from_channel
 from playwright._impl._fetch import APIRequest
 from playwright._impl._local_utils import LocalUtils
-from playwright._impl._selectors import Selectors
+from playwright._impl._selectors import Selectors, SelectorsOwner
 
 
 class Playwright(ChannelOwner):
@@ -40,7 +40,14 @@ def __init__(
         self.firefox._playwright = self
         self.webkit = from_channel(initializer["webkit"])
         self.webkit._playwright = self
-        self.selectors = from_channel(initializer["selectors"])
+
+        self.selectors = Selectors(self._loop)
+        selectors_owner: SelectorsOwner = from_channel(initializer["selectors"])
+        self.selectors._add_channel(selectors_owner)
+
+        self._connection.on(
+            "close", lambda: self.selectors._remove_channel(selectors_owner)
+        )
         self.devices = {}
         self.devices = {
             device["name"]: parse_device_descriptor(device["descriptor"])
@@ -57,6 +64,12 @@ def __getitem__(self, value: str) -> "BrowserType":
             return self.webkit
         raise ValueError("Invalid browser " + value)
 
+    def _set_selectors(self, selectors: SelectorsOwner) -> None:
+        selectors_owner = from_channel(self._initializer["selectors"])
+        self.selectors._remove_channel(selectors_owner)
+        self.selectors = selectors
+        self.selectors._add_channel(selectors_owner)
+
     def stop(self) -> None:
         pass
 

playwright/_impl/_selectors.py

@@ -12,19 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 from pathlib import Path
-from typing import Any, Dict, Union
+from typing import Any, Dict, List, Set, Union
 
 from playwright._impl._api_types import Error
 from playwright._impl._connection import ChannelOwner
 from playwright._impl._helper import async_readfile
 
 
-class Selectors(ChannelOwner):
-    def __init__(
-        self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
-    ) -> None:
-        super().__init__(parent, type, guid, initializer)
+class Selectors:
+    def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
+        self._loop = loop
+        self._channels: Set[SelectorsOwner] = set()
+        self._registrations: List[Dict] = []
 
     async def register(
         self,
@@ -40,4 +41,23 @@ async def register(
         params: Dict[str, Any] = dict(name=name, source=script)
         if contentScript:
             params["contentScript"] = True
-        await self._channel.send("register", params)
+        for channel in self._channels:
+            await channel._channel.send("register", params)
+        self._registrations.append(params)
+
+    def _add_channel(self, channel: "SelectorsOwner") -> None:
+        self._channels.add(channel)
+        for params in self._registrations:
+            # This should not fail except for connection closure, but just in case we catch.
+            channel._channel.send_no_reply("register", params)
+
+    def _remove_channel(self, channel: "SelectorsOwner") -> None:
+        if channel in self._channels:
+            self._channels.remove(channel)
+
+
+class SelectorsOwner(ChannelOwner):
+    def __init__(
+        self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
+    ) -> None:
+        super().__init__(parent, type, guid, initializer)