feat(expect): add types expectations into the page & browser context. (#88)

Author: pavelfeldman

client.py

@@ -22,9 +22,9 @@ def main() -> None:
         "<button id=button onclick=\"window.open('http://webkit.org', '_blank')\">Click me</input>"
     )
 
-    with page.expect_event("popup") as popup:
+    with page.expect_popup() as popup_info:
         page.click("#button")
-    print(popup.value)
+    print(popup_info.value)
 
     print("Contexts in browser: %d" % len(browser.contexts))
     print("Creating context...")

playwright/browser_context.py

@@ -17,6 +17,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from playwright.connection import ChannelOwner, ConnectionScope, from_channel
+from playwright.event_context_manager import AsyncEventContextManager
 from playwright.helper import (
     Cookie,
     Error,
@@ -212,3 +213,13 @@ async def close(self) -> None:
             return
         self._is_closed_or_closing = True
         await self._channel.send("close")
+
+    def expect_event(
+        self, event: str, predicate: Callable[[Any], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager:
+        return AsyncEventContextManager(self, event, predicate, timeout)
+
+    def expect_page(
+        self, predicate: Callable[[Page], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[Page]:
+        return AsyncEventContextManager(self, "page", predicate, timeout)

playwright/event_context_manager.py

@@ -0,0 +1,62 @@
+# 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
+from typing import Any, Callable, Generic, Optional, TypeVar, cast
+
+from playwright.connection import ChannelOwner
+from playwright.wait_helper import WaitHelper
+
+T = TypeVar("T")
+
+
+class AsyncEventInfo(Generic[T]):
+    def __init__(
+        self,
+        channel_owner: ChannelOwner,
+        event: str,
+        predicate: Callable[[T], bool] = None,
+        timeout: int = None,
+    ) -> None:
+        self._value: Optional[T] = None
+        wait_helper = WaitHelper()
+        wait_helper.reject_on_timeout(
+            timeout or 30000, f'Timeout while waiting for event "${event}"'
+        )
+        self._future = asyncio.get_event_loop().create_task(
+            wait_helper.wait_for_event(channel_owner, event, predicate)
+        )
+
+    @property
+    async def value(self) -> T:
+        if not self._value:
+            self._value = await self._future
+        return cast(T, self._value)
+
+
+class AsyncEventContextManager(Generic[T]):
+    def __init__(
+        self,
+        channel_owner: ChannelOwner,
+        event: str,
+        predicate: Callable[[T], bool] = None,
+        timeout: int = None,
+    ) -> None:
+        self._event = AsyncEventInfo(channel_owner, event, predicate, timeout)
+
+    async def __aenter__(self) -> AsyncEventInfo[T]:
+        return self._event
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        await self._event.value

playwright/file_chooser.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, List, Union
+from typing import TYPE_CHECKING, Any, List, Union
 
 from playwright.helper import FilePayload
 
@@ -25,6 +25,7 @@ class FileChooser:
     def __init__(
         self, page: "Page", element_handle: "ElementHandle", is_multiple: bool
     ) -> None:
+        self._sync_owner: Any = None
         self._page = page
         self._element_handle = element_handle
         self._is_multiple = is_multiple

playwright/page.py

@@ -25,7 +25,11 @@
     from_channel,
     from_nullable_channel,
 )
+from playwright.console_message import ConsoleMessage
+from playwright.dialog import Dialog
+from playwright.download import Download
 from playwright.element_handle import ElementHandle, ValuesToSelect
+from playwright.event_context_manager import AsyncEventContextManager
 from playwright.file_chooser import FileChooser
 from playwright.frame import Frame
 from playwright.helper import (
@@ -99,7 +103,7 @@ def __init__(self, scope: ConnectionScope, guid: str, initializer: Dict) -> None
         self._main_frame: Frame = from_channel(initializer["mainFrame"])
         self._main_frame._page = self
         self._frames = [self._main_frame]
-        self._viewport_size = initializer["viewportSize"]
+        self._viewport_size = initializer.get("viewportSize")
         self._is_closed = False
         self._workers: List[Worker] = list()
         self._bindings: Dict[str, Any] = dict()
@@ -704,6 +708,51 @@ async def pdf(
         binary = await self._channel.send("pdf", locals_to_params(locals()))
         return base64.b64decode(binary)
 
+    def expect_event(
+        self, event: str, predicate: Callable[[Any], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager:
+        return AsyncEventContextManager(self, event, predicate, timeout)
+
+    def expect_console_message(
+        self, predicate: Callable[[ConsoleMessage], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[ConsoleMessage]:
+        return AsyncEventContextManager(self, "console", predicate, timeout)
+
+    def expect_dialog(
+        self, predicate: Callable[[Dialog], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[Dialog]:
+        return AsyncEventContextManager(self, "dialog", predicate, timeout)
+
+    def expect_download(
+        self, predicate: Callable[[Download], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[Download]:
+        return AsyncEventContextManager(self, "download", predicate, timeout)
+
+    def expect_file_chooser(
+        self, predicate: Callable[[FileChooser], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[FileChooser]:
+        return AsyncEventContextManager(self, "filechooser", predicate, timeout)
+
+    def expect_request(
+        self, predicate: Callable[[Request], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[Request]:
+        return AsyncEventContextManager(self, "request", predicate, timeout)
+
+    def expect_response(
+        self, predicate: Callable[[Response], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[Response]:
+        return AsyncEventContextManager(self, "response", predicate, timeout)
+
+    def expect_popup(
+        self, predicate: Callable[["Page"], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager["Page"]:
+        return AsyncEventContextManager(self, "popup", predicate, timeout)
+
+    def expect_worker(
+        self, predicate: Callable[[Worker], bool] = None, timeout: int = None,
+    ) -> AsyncEventContextManager[Worker]:
+        return AsyncEventContextManager(self, "worker", predicate, timeout)
+
 
 class BindingCall(ChannelOwner):
     def __init__(self, scope: ConnectionScope, guid: str, initializer: Dict) -> None:

playwright/sync.py

@@ -15,7 +15,7 @@
 import sys
 import typing
 
-from playwright.sync_base import SyncBase, mapping
+from playwright.sync_base import EventContextManager, SyncBase, mapping
 
 if sys.version_info >= (3, 8):  # pragma: no cover
     from typing import Literal
@@ -31,6 +31,7 @@
 from playwright.dialog import Dialog as DialogAsync
 from playwright.download import Download as DownloadAsync
 from playwright.element_handle import ElementHandle as ElementHandleAsync
+from playwright.file_chooser import FileChooser as FileChooserAsync
 from playwright.frame import Frame as FrameAsync
 from playwright.helper import (
     ConsoleMessageLocation,
@@ -770,6 +771,67 @@ def snapshot(
 mapping.register(AccessibilityAsync, Accessibility)
 
 
+class FileChooser(SyncBase):
+    def __init__(self, obj: FileChooserAsync):
+        super().__init__(obj)
+
+    def as_async(self) -> FileChooserAsync:
+        return self._async_obj
+
+    @classmethod
+    def _from_async(cls, obj: FileChooserAsync) -> "FileChooser":
+        if not obj._sync_owner:
+            obj._sync_owner = cls(obj)
+        return obj._sync_owner
+
+    @classmethod
+    def _from_async_nullable(
+        cls, obj: FileChooserAsync = None
+    ) -> typing.Optional["FileChooser"]:
+        return FileChooser._from_async(obj) if obj else None
+
+    @classmethod
+    def _from_async_list(
+        cls, items: typing.List[FileChooserAsync]
+    ) -> typing.List["FileChooser"]:
+        return list(map(lambda a: FileChooser._from_async(a), items))
+
+    @classmethod
+    def _from_async_dict(
+        cls, map: typing.Dict[str, FileChooserAsync]
+    ) -> typing.Dict[str, "FileChooser"]:
+        return {name: FileChooser._from_async(value) for name, value in map.items()}
+
+    @property
+    def page(self) -> "Page":
+        return Page._from_async(self._async_obj.page)
+
+    @property
+    def element(self) -> "ElementHandle":
+        return ElementHandle._from_async(self._async_obj.element)
+
+    @property
+    def isMultiple(self) -> bool:
+        return self._async_obj.isMultiple
+
+    def setFiles(
+        self,
+        files: typing.Union[
+            str, FilePayload, typing.List[str], typing.List[FilePayload]
+        ],
+        timeout: int = None,
+        noWaitAfter: bool = None,
+    ) -> NoneType:
+        return self._sync(
+            self._async_obj.setFiles(
+                files=files, timeout=timeout, noWaitAfter=noWaitAfter
+            )
+        )
+
+
+mapping.register(FileChooserAsync, FileChooser)
+
+
 class Frame(SyncBase):
     def __init__(self, obj: FrameAsync):
         super().__init__(obj)
@@ -2119,6 +2181,62 @@ def pdf(
             )
         )
 
+    def expect_console_message(
+        self,
+        predicate: typing.Union[typing.Callable[["ConsoleMessage"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["ConsoleMessage"]:
+        return EventContextManager(self, "console", predicate, timeout)
+
+    def expect_dialog(
+        self,
+        predicate: typing.Union[typing.Callable[["Dialog"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["Dialog"]:
+        return EventContextManager(self, "dialog", predicate, timeout)
+
+    def expect_download(
+        self,
+        predicate: typing.Union[typing.Callable[["Download"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["Download"]:
+        return EventContextManager(self, "download", predicate, timeout)
+
+    def expect_file_chooser(
+        self,
+        predicate: typing.Union[typing.Callable[["FileChooser"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["FileChooser"]:
+        return EventContextManager(self, "filechooser", predicate, timeout)
+
+    def expect_request(
+        self,
+        predicate: typing.Union[typing.Callable[["Request"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["Request"]:
+        return EventContextManager(self, "request", predicate, timeout)
+
+    def expect_response(
+        self,
+        predicate: typing.Union[typing.Callable[["Response"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["Response"]:
+        return EventContextManager(self, "response", predicate, timeout)
+
+    def expect_popup(
+        self,
+        predicate: typing.Union[typing.Callable[["Page"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["Page"]:
+        return EventContextManager(self, "popup", predicate, timeout)
+
+    def expect_worker(
+        self,
+        predicate: typing.Union[typing.Callable[["Worker"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["Worker"]:
+        return EventContextManager(self, "worker", predicate, timeout)
+
 
 mapping.register(PageAsync, Page)
 
@@ -2244,6 +2362,13 @@ def waitForEvent(
     def close(self) -> NoneType:
         return self._sync(self._async_obj.close())
 
+    def expect_page(
+        self,
+        predicate: typing.Union[typing.Callable[["Page"], bool]] = None,
+        timeout: int = None,
+    ) -> EventContextManager["Page"]:
+        return EventContextManager(self, "page", predicate, timeout)
+
 
 mapping.register(BrowserContextAsync, BrowserContext)