Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 14 additions & 8 deletions playwright/_impl/_browser_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
from playwright._impl._console_message import ConsoleMessage
from playwright._impl._debugger import Debugger
from playwright._impl._dialog import Dialog
from playwright._impl._disposable import Disposable, DisposableStub
from playwright._impl._errors import Error, TargetClosedError
from playwright._impl._event_context_manager import EventContextManagerImpl
from playwright._impl._fetch import APIRequestContext
Expand Down Expand Up @@ -394,16 +395,18 @@ async def set_offline(self, offline: bool) -> None:

async def add_init_script(
self, script: str = None, path: Union[str, Path] = None
) -> None:
) -> Disposable:
if path:
script = (await async_readfile(path)).decode()
if not isinstance(script, str):
raise Error("Either path or script parameter must be specified")
await self._channel.send("addInitScript", None, dict(source=script))
return from_channel(
await self._channel.send("addInitScript", None, dict(source=script))
)

async def expose_binding(
self, name: str, callback: Callable, handle: bool = None
) -> None:
) -> Disposable:
for page in self._pages:
if name in page._bindings:
raise Error(
Expand All @@ -412,16 +415,18 @@ async def expose_binding(
if name in self._bindings:
raise Error(f'Function "{name}" has been already registered')
self._bindings[name] = callback
await self._channel.send(
"exposeBinding", None, dict(name=name, needsHandle=handle or False)
return from_channel(
await self._channel.send(
"exposeBinding", None, dict(name=name, needsHandle=handle or False)
)
)

async def expose_function(self, name: str, callback: Callable) -> None:
await self.expose_binding(name, lambda source, *args: callback(*args))
async def expose_function(self, name: str, callback: Callable) -> Disposable:
return await self.expose_binding(name, lambda source, *args: callback(*args))

async def route(
self, url: URLMatch, handler: RouteHandlerCallback, times: int = None
) -> None:
) -> DisposableStub:
self._routes.insert(
0,
RouteHandler(
Expand All @@ -433,6 +438,7 @@ async def route(
),
)
await self._update_interception_patterns()
return DisposableStub(lambda: self.unroute(url, handler), self)

async def unroute(
self, url: URLMatch, handler: Optional[RouteHandlerCallback] = None
Expand Down
93 changes: 93 additions & 0 deletions playwright/_impl/_disposable.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
# Copyright (c) Microsoft Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import inspect
import traceback
from typing import Awaitable, Callable, Dict

import greenlet

from playwright._impl._connection import ChannelOwner
from playwright._impl._errors import Error, is_target_closed_error


class Disposable(ChannelOwner):
def __init__(
self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
) -> None:
super().__init__(parent, type, guid, initializer)

async def dispose(self) -> None:
try:
await self._channel.send(
"dispose",
None,
)
except Exception as e:
if not is_target_closed_error(e):
raise e

async def close(self) -> None:
await self.dispose()

def __repr__(self) -> str:
return "<Disposable>"


class DisposableStub:
def __init__(
self,
dispose_fn: Callable[[], Awaitable[None]],
parent: ChannelOwner,
) -> None:
self._dispose_fn = dispose_fn
self._loop = parent._loop
self._dispatcher_fiber = parent._dispatcher_fiber

async def dispose(self) -> None:
await self._dispose_fn()

async def __aenter__(self) -> "DisposableStub":
return self

async def __aexit__(self, *args: object) -> None:
await self.dispose()

def __enter__(self) -> "DisposableStub":
return self

def __exit__(self, *args: object) -> None:
self._sync(self.dispose())

def _sync(self, coro: object) -> object:
__tracebackhide__ = True
if self._loop.is_closed():
coro.close() # type: ignore
raise Error("Event loop is closed! Is Playwright already stopped?")
g_self = greenlet.getcurrent()
task = self._loop.create_task(coro) # type: ignore
setattr(task, "__pw_stack__", inspect.stack(0))
setattr(task, "__pw_stack_trace__", traceback.extract_stack(limit=10))
task.add_done_callback(lambda _: g_self.switch())
while not task.done():
self._dispatcher_fiber.switch() # type: ignore
asyncio._set_running_loop(self._loop)
return task.result()

async def close(self) -> None:
await self.dispose()

def __repr__(self) -> str:
return "<Disposable>"
3 changes: 3 additions & 0 deletions playwright/_impl/_object_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from playwright._impl._connection import ChannelOwner
from playwright._impl._debugger import Debugger
from playwright._impl._dialog import Dialog
from playwright._impl._disposable import Disposable
from playwright._impl._element_handle import ElementHandle
from playwright._impl._fetch import APIRequestContext
from playwright._impl._frame import Frame
Expand Down Expand Up @@ -69,6 +70,8 @@ def create_remote_object(
return Debugger(parent, type, guid, initializer)
if type == "Dialog":
return Dialog(parent, type, guid, initializer)
if type == "Disposable":
return Disposable(parent, type, guid, initializer)
if type == "ElementHandle":
return ElementHandle(parent, type, guid, initializer)
if type == "Frame":
Expand Down
26 changes: 16 additions & 10 deletions playwright/_impl/_page.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
from_nullable_channel,
)
from playwright._impl._console_message import ConsoleMessage
from playwright._impl._disposable import Disposable, DisposableStub
from playwright._impl._download import Download
from playwright._impl._element_handle import ElementHandle, determine_screenshot_type
from playwright._impl._errors import Error, TargetClosedError, is_target_closed_error
Expand Down Expand Up @@ -501,23 +502,25 @@ async def add_style_tag(
) -> ElementHandle:
return await self._main_frame.add_style_tag(**locals_to_params(locals()))

async def expose_function(self, name: str, callback: Callable) -> None:
await self.expose_binding(name, lambda source, *args: callback(*args))
async def expose_function(self, name: str, callback: Callable) -> Disposable:
return await self.expose_binding(name, lambda source, *args: callback(*args))

async def expose_binding(
self, name: str, callback: Callable, handle: bool = None
) -> None:
) -> Disposable:
if name in self._bindings:
raise Error(f'Function "{name}" has been already registered')
if name in self._browser_context._bindings:
raise Error(
f'Function "{name}" has been already registered in the browser context'
)
self._bindings[name] = callback
await self._channel.send(
"exposeBinding",
None,
dict(name=name, needsHandle=handle or False),
return from_channel(
await self._channel.send(
"exposeBinding",
None,
dict(name=name, needsHandle=handle or False),
)
)

async def set_extra_http_headers(self, headers: Dict[str, str]) -> None:
Expand Down Expand Up @@ -661,18 +664,20 @@ async def bring_to_front(self) -> None:

async def add_init_script(
self, script: str = None, path: Union[str, Path] = None
) -> None:
) -> Disposable:
if path:
script = add_source_url_to_script(
(await async_readfile(path)).decode(), path
)
if not isinstance(script, str):
raise Error("Either path or script parameter must be specified")
await self._channel.send("addInitScript", None, dict(source=script))
return from_channel(
await self._channel.send("addInitScript", None, dict(source=script))
)

async def route(
self, url: URLMatch, handler: RouteHandlerCallback, times: int = None
) -> None:
) -> DisposableStub:
self._routes.insert(
0,
RouteHandler(
Expand All @@ -684,6 +689,7 @@ async def route(
),
)
await self._update_interception_patterns()
return DisposableStub(lambda: self.unroute(url, handler), self)

async def unroute(
self, url: URLMatch, handler: Optional[RouteHandlerCallback] = None
Expand Down
18 changes: 14 additions & 4 deletions playwright/_impl/_screencast.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from playwright._impl._api_structures import ScreencastFrame
from playwright._impl._artifact import Artifact
from playwright._impl._connection import from_nullable_channel
from playwright._impl._disposable import DisposableStub
from playwright._impl._errors import Error
from playwright._impl._helper import locals_to_params

Expand Down Expand Up @@ -63,7 +64,7 @@ async def start(
onFrame: ScreencastFrameCallback = None,
path: Union[str, Path] = None,
quality: int = None,
) -> None:
) -> DisposableStub:
if self._started:
raise Error("Screencast is already started")
self._started = True
Expand All @@ -81,6 +82,7 @@ async def start(
if artifact_channel:
self._artifact = from_nullable_channel(artifact_channel)
self._save_path = path
return DisposableStub(lambda: self.stop(), self._page)

async def stop(self) -> None:
self._started = False
Expand All @@ -96,18 +98,26 @@ async def show_actions(
duration: float = None,
position: ScreencastPosition = None,
fontSize: int = None,
) -> None:
) -> DisposableStub:
await self._page._channel.send(
"screencastShowActions", None, locals_to_params(locals())
)
return DisposableStub(lambda: self.hide_actions(), self._page)

async def hide_actions(self) -> None:
await self._page._channel.send("screencastHideActions", None)

async def show_overlay(self, html: str, duration: float = None) -> None:
await self._page._channel.send(
async def show_overlay(self, html: str, duration: float = None) -> DisposableStub:
result = await self._page._channel.send_return_as_dict(
"screencastShowOverlay", None, locals_to_params(locals())
)
overlay_id = (result or {}).get("id")
return DisposableStub(
lambda: self._page._channel.send(
"screencastRemoveOverlay", None, {"id": overlay_id}
),
self._page,
)

async def show_chapter(
self,
Expand Down
6 changes: 5 additions & 1 deletion playwright/_impl/_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from playwright._impl._api_structures import TracingGroupLocation
from playwright._impl._artifact import Artifact
from playwright._impl._connection import ChannelOwner, from_nullable_channel
from playwright._impl._disposable import DisposableStub
from playwright._impl._helper import locals_to_params


Expand Down Expand Up @@ -148,8 +149,11 @@ def _reset_stack_counter(self) -> None:
self._is_tracing = False
self._connection.set_is_tracing(False)

async def group(self, name: str, location: TracingGroupLocation = None) -> None:
async def group(
self, name: str, location: TracingGroupLocation = None
) -> DisposableStub:
await self._channel.send("tracingGroup", None, locals_to_params(locals()))
return DisposableStub(lambda: self.group_end(), self)

async def group_end(self) -> None:
await self._channel.send(
Expand Down
Loading
Loading