class PlaywrightContextManager:
def __init__(self) -> None:
self._connection: Connection
async def __aenter__(self) -> AsyncPlaywright:
self._connection = Connection(
None,
create_remote_object,
PipeTransport(asyncio.get_event_loop(), compute_driver_executable()),
)
loop = asyncio.get_running_loop()
self._connection._loop = loop
loop.create_task(self._connection.run())
playwright_future = asyncio.create_task(
self._connection.wait_for_object_with_known_name("Playwright")
)
done, pending = await asyncio.wait(
{self._connection._transport.on_error_future, playwright_future},
return_when=asyncio.FIRST_COMPLETED,
)
if not playwright_future.done():
playwright_future.cancel()
playwright = AsyncPlaywright(next(iter(done)).result())
playwright.stop = self.__aexit__ # type: ignore
return playwright
async def start(self) -> AsyncPlaywright:
return await self.__aenter__()
async def __aexit__(self, *args: Any) -> None:
await self._connection.stop_async()
async def connect(
self,
ws_endpoint: str,
timeout: float = None,
slow_mo: float = None,
headers: Dict[str, str] = None,
) -> Browser:
if timeout is None:
timeout = 30000
transport = WebSocketTransport(self._connection._loop, ws_endpoint,
headers, slow_mo)
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())
future = connection._loop.create_task(
connection.wait_for_object_with_known_name("Playwright"))
timeout_future = throw_on_timeout(timeout,
Error("Connection timed out"))
done, pending = await asyncio.wait(
{transport.on_error_future, future, timeout_future},
return_when=asyncio.FIRST_COMPLETED,
)
if not future.done():
future.cancel()
if not timeout_future.done():
timeout_future.cancel()
playwright = next(iter(done)).result()
self._connection._child_ws_connections.append(connection)
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
def handle_transport_close() -> None:
for context in browser.contexts:
for page in context.pages:
page._on_close()
context._on_close()
browser._on_close()
transport.once("close", handle_transport_close)
return browser
async def connect(
self,
ws_endpoint: str,
timeout: float = None,
slow_mo: float = None,
headers: Dict[str, str] = None,
) -> Browser:
transport = WebSocketTransport(ws_endpoint, timeout, headers)
connection = Connection(
self._connection._dispatcher_fiber,
self._connection._object_factory,
transport,
)
await connection._transport.start()
connection._is_sync = self._connection._is_sync
connection._loop = self._connection._loop
connection._loop.create_task(connection.run())
obj = asyncio.create_task(
connection.wait_for_object_with_known_name("Playwright"))
done, pending = await asyncio.wait(
{
obj,
connection._transport.on_error_future, # type: ignore
},
return_when=asyncio.FIRST_COMPLETED,
)
if not obj.done():
obj.cancel()
playwright = next(iter(done)).result()
self._connection._child_ws_connections.append(connection)
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