feat: add AsyncContext/Disposable support (#3052)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: Skn0tt

playwright/_impl/_browser_context.py

@@ -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
@@ -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(
@@ -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(
@@ -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

playwright/_impl/_disposable.py

@@ -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>"

playwright/_impl/_object_factory.py

@@ -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
@@ -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":

playwright/_impl/_page.py

@@ -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
@@ -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:
@@ -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(
@@ -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

playwright/_impl/_screencast.py

@@ -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
 
@@ -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
@@ -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
@@ -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,

playwright/_impl/_tracing.py

@@ -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
 
 
@@ -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(