From aac5fe7e1ac6669888bc28a0904c45938690e3d6 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Thu, 23 Jul 2020 21:12:09 -0700
Subject: [PATCH] feat(expect): add types expectations into the page & browser
 context.

---
 client.py                           |   4 +-
 playwright/browser_context.py       |  11 +++
 playwright/event_context_manager.py |  62 ++++++++++++++
 playwright/file_chooser.py          |   3 +-
 playwright/page.py                  |  51 ++++++++++-
 playwright/sync.py                  | 127 +++++++++++++++++++++++++++-
 playwright/sync_base.py             |  27 +++---
 scripts/generate_sync_api.py        |  26 +++++-
 tests/test_browsercontext.py        |  27 +++---
 tests/test_page.py                  |  20 ++---
 10 files changed, 310 insertions(+), 48 deletions(-)
 create mode 100644 playwright/event_context_manager.py

diff --git a/client.py b/client.py
index 30277ae5c..1fc7dc19b 100644
--- a/client.py
+++ b/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...")
diff --git a/playwright/browser_context.py b/playwright/browser_context.py
index dff8d4415..cbb13b737 100644
--- a/playwright/browser_context.py
+++ b/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)
diff --git a/playwright/event_context_manager.py b/playwright/event_context_manager.py
new file mode 100644
index 000000000..7ea6e5ef9
--- /dev/null
+++ b/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
diff --git a/playwright/file_chooser.py b/playwright/file_chooser.py
index 4ca895cad..b4d224918 100644
--- a/playwright/file_chooser.py
+++ b/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
diff --git a/playwright/page.py b/playwright/page.py
index 96b17368c..4162d1c34 100644
--- a/playwright/page.py
+++ b/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:
diff --git a/playwright/sync.py b/playwright/sync.py
index 7f66a36da..14e81c8b0 100644
--- a/playwright/sync.py
+++ b/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)
 
diff --git a/playwright/sync_base.py b/playwright/sync_base.py
index bff56e805..ccde5f646 100644
--- a/playwright/sync_base.py
+++ b/playwright/sync_base.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import asyncio
-from typing import Any, Callable, List, Tuple, Union
+from typing import Any, Callable, Generic, List, Optional, Tuple, TypeVar, cast
 
 from playwright.wait_helper import WaitHelper
 
@@ -35,16 +35,18 @@ def get_sync_class(self, input_async_inst: object) -> Any:
 
 mapping = AsyncToSyncMapping()
 
+T = TypeVar("T")
 
-class Event:
+
+class EventInfo(Generic[T]):
     def __init__(
         self,
         sync_base: "SyncBase",
         event: str,
-        predicate: Callable[[Any], bool] = None,
+        predicate: Callable[[T], bool] = None,
         timeout: int = None,
     ) -> None:
-        self._value: Any = None
+        self._value: Optional[T] = None
 
         wait_helper = WaitHelper()
         wait_helper.reject_on_timeout(
@@ -55,24 +57,24 @@ def __init__(
         )
 
     @property
-    def value(self) -> Any:
+    def value(self) -> T:
         if not self._value:
             value = loop.run_until_complete(self._future)
             self._value = mapping.get_sync_class(value)._from_async(value)
-        return self._value
+        return cast(T, self._value)
 
 
-class EventContextManager:
+class EventContextManager(Generic[T]):
     def __init__(
         self,
         sync_base: "SyncBase",
         event: str,
-        predicate: Callable[[Any], bool] = None,
+        predicate: Callable[[T], bool] = None,
         timeout: int = None,
     ) -> None:
-        self._event = Event(sync_base, event, predicate, timeout)
+        self._event = EventInfo(sync_base, event, predicate, timeout)
 
-    def __enter__(self) -> Event:
+    def __enter__(self) -> EventInfo[T]:
         return self._event
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
@@ -102,9 +104,6 @@ def remove_listener(self, event_name: str, handler: Any) -> None:
         self._async_obj.remove_listener(event_name, handler)
 
     def expect_event(
-        self,
-        event: str,
-        predicate: Union[Callable[[Any], bool]] = None,
-        timeout: int = None,
+        self, event: str, predicate: Callable[[Any], bool] = None, timeout: int = None,
     ) -> EventContextManager:
         return EventContextManager(self, event, predicate, timeout)
diff --git a/scripts/generate_sync_api.py b/scripts/generate_sync_api.py
index 05d99f795..a76d0d3d3 100644
--- a/scripts/generate_sync_api.py
+++ b/scripts/generate_sync_api.py
@@ -33,6 +33,7 @@
 from playwright.dialog import Dialog
 from playwright.download import Download
 from playwright.element_handle import ElementHandle, ValuesToSelect
+from playwright.file_chooser import FileChooser
 from playwright.frame import Frame
 from playwright.input import Keyboard, Mouse
 from playwright.js_handle import JSHandle
@@ -191,7 +192,11 @@ def generate(t: Any) -> None:
             prefix = "        return " + prefix + f"self._async_obj.{name}"
             print(f"{prefix}{arguments(value, len(prefix))}{suffix}")
     for [name, value] in t.__dict__.items():
-        if not name.startswith("_") and isinstance(value, FunctionType):
+        if (
+            not name.startswith("_")
+            and isinstance(value, FunctionType)
+            and "expect_" not in name
+        ):
             print("")
             print(
                 f"    def {name}({signature(value, len(name) + 9)}) -> {return_type(value)}:"
@@ -200,6 +205,21 @@ def generate(t: Any) -> None:
             prefix = "        return " + prefix + f"self._sync(self._async_obj.{name}("
             suffix = "))" + suffix
             print(f"{prefix}{arguments(value, len(prefix))}{suffix}")
+        if "expect_" in name and "expect_event" not in name:
+            print("")
+            return_type_value = return_type(value)
+            return_type_value = re.sub(r"\"Async([^\"]+)\"", r"\1", return_type_value)
+            event_name = re.sub(r"expect_(.*)", r"\1", name)
+            event_name = re.sub(r"_", "", event_name)
+            event_name = re.sub(r"consolemessage", "console", event_name)
+            print(
+                f"    def {name}({signature(value, len(name) + 9)}) -> {return_type_value}:"
+            )
+            print(
+                f'        return EventContextManager(self, "{event_name}", predicate, timeout)'
+            )
+
+    print("")
     print(f"mapping.register({short_name(t)}Async, {short_name(t)})")
 
 
@@ -223,7 +243,7 @@ def main() -> None:
 
 import typing
 import sys
-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
@@ -239,6 +259,7 @@ def main() -> None:
 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, Error, FilePayload, SelectOption, Viewport
 from playwright.input import Keyboard as KeyboardAsync, Mouse as MouseAsync
@@ -261,6 +282,7 @@ def main() -> None:
         JSHandle,
         ElementHandle,
         Accessibility,
+        FileChooser,
         Frame,
         Worker,
         Selectors,
diff --git a/tests/test_browsercontext.py b/tests/test_browsercontext.py
index 800c9ee74..aebc77c84 100644
--- a/tests/test_browsercontext.py
+++ b/tests/test_browsercontext.py
@@ -577,19 +577,17 @@ async def test_offline_should_emulate_navigator_online(context, server):
 
 async def test_page_event_should_have_url(context, server):
     page = await context.newPage()
-    [other_page, _] = await asyncio.gather(
-        context.waitForEvent("page"),
-        page.evaluate("url => window.open(url)", server.EMPTY_PAGE),
-    )
+    async with context.expect_page() as other_page_info:
+        await page.evaluate("url => window.open(url)", server.EMPTY_PAGE)
+    other_page = await other_page_info.value
     assert other_page.url == server.EMPTY_PAGE
 
 
 async def test_page_event_should_have_url_after_domcontentloaded(context, server):
     page = await context.newPage()
-    [other_page, _] = await asyncio.gather(
-        context.waitForEvent("page"),
-        page.evaluate("url => window.open(url)", server.EMPTY_PAGE),
-    )
+    async with context.expect_page() as other_page_info:
+        await page.evaluate("url => window.open(url)", server.EMPTY_PAGE)
+    other_page = await other_page_info.value
     await other_page.waitForLoadState("domcontentloaded")
     assert other_page.url == server.EMPTY_PAGE
 
@@ -598,10 +596,9 @@ async def test_page_event_should_have_about_blank_url_with_domcontentloaded(
     context, server
 ):
     page = await context.newPage()
-    [other_page, _] = await asyncio.gather(
-        context.waitForEvent("page"),
-        page.evaluate("url => window.open(url)", "about:blank"),
-    )
+    async with context.expect_page() as other_page_info:
+        await page.evaluate("url => window.open(url)", "about:blank")
+    other_page = await other_page_info.value
     await other_page.waitForLoadState("domcontentloaded")
     assert other_page.url == "about:blank"
 
@@ -610,9 +607,9 @@ async def test_page_event_should_have_about_blank_for_empty_url_with_domcontentl
     context, server
 ):
     page = await context.newPage()
-    [other_page, _] = await asyncio.gather(
-        context.waitForEvent("page"), page.evaluate("window.open()")
-    )
+    async with context.expect_page() as other_page_info:
+        await page.evaluate("window.open()")
+    other_page = await other_page_info.value
     await other_page.waitForLoadState("domcontentloaded")
     assert other_page.url == "about:blank"
 
diff --git a/tests/test_page.py b/tests/test_page.py
index 8964c1a26..cf22c57b8 100644
--- a/tests/test_page.py
+++ b/tests/test_page.py
@@ -1186,22 +1186,18 @@ async def test_fill_should_be_able_to_clear(page, server):
 
 
 async def test_close_event_should_work_with_window_close(page, server):
-    promise = asyncio.create_task(page.waitForEvent("popup"))
-    await asyncio.sleep(0)  # execute scheduled tasks, but don't await them
-    await page.evaluate("window['newPage'] = window.open('about:blank')")
-    popup = await promise
-    close_promise = asyncio.create_task(popup.waitForEvent("close"))
-    await asyncio.sleep(0)  # execute scheduled tasks, but don't await them
-    await page.evaluate("window['newPage'].close()")
-    await close_promise
+    async with page.expect_popup() as popup_info:
+        await page.evaluate("window['newPage'] = window.open('about:blank')")
+    popup = await popup_info.value
+
+    async with popup.expect_event("close"):
+        await page.evaluate("window['newPage'].close()")
 
 
 async def test_close_event_should_work_with_page_close(context, server):
     page = await context.newPage()
-    close_promise = asyncio.create_task(page.waitForEvent("close"))
-    await asyncio.sleep(0)  # execute scheduled tasks, but don't await them
-    await page.close()
-    await close_promise
+    async with page.expect_event("close"):
+        await page.close()
 
 
 async def test_page_context_should_return_the_correct_browser_instance(page, context):