fix: improve typing for SyncBase and AsyncBase (microsoft#790)

kumaraditya303 · web-flow · commit 673eed8f30af · 2021-07-15T11:37:26.000-07:00
diff --git a/playwright/_impl/_async_base.py b/playwright/_impl/_async_base.py
@@ -15,19 +15,19 @@
 import asyncio
 import traceback
 from types import TracebackType
-from typing import Any, Awaitable, Callable, Generic, Type, TypeVar
+from typing import Any, Awaitable, Callable, Generic, Type, TypeVar, Union
 
 from playwright._impl._impl_to_api_mapping import ImplToApiMapping, ImplWrapper
 
 mapping = ImplToApiMapping()
 
 
 T = TypeVar("T")
-Self = TypeVar("Self", bound="AsyncBase")
+Self = TypeVar("Self", bound="AsyncContextManager")
 
 
 class AsyncEventInfo(Generic[T]):
-    def __init__(self, future: asyncio.Future) -> None:
+    def __init__(self, future: "asyncio.Future[T]") -> None:
         self._future = future
 
     @property
@@ -39,13 +39,18 @@ def is_done(self) -> bool:
 
 
 class AsyncEventContextManager(Generic[T]):
-    def __init__(self, future: asyncio.Future) -> None:
-        self._event: AsyncEventInfo = AsyncEventInfo(future)
+    def __init__(self, future: "asyncio.Future[T]") -> None:
+        self._event = AsyncEventInfo[T](future)
 
     async def __aenter__(self) -> AsyncEventInfo[T]:
         return self._event
 
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+    async def __aexit__(
+        self,
+        exc_type: Type[BaseException],
+        exc_val: BaseException,
+        exc_tb: TracebackType,
+    ) -> None:
         await self._event.value
 
 
@@ -68,17 +73,19 @@ def _wrap_handler(self, handler: Any) -> Callable[..., None]:
             return mapping.wrap_handler(handler)
         return handler
 
-    def on(self, event: str, f: Any) -> None:
+    def on(self, event: str, f: Callable[..., Union[Awaitable[None], None]]) -> None:
         """Registers the function ``f`` to the event name ``event``."""
         self._impl_obj.on(event, self._wrap_handler(f))
 
-    def once(self, event: str, f: Any) -> None:
+    def once(self, event: str, f: Callable[..., Union[Awaitable[None], None]]) -> None:
         """The same as ``self.on``, except that the listener is automatically
         removed after being called.
         """
         self._impl_obj.once(event, self._wrap_handler(f))
 
-    def remove_listener(self, event: str, f: Any) -> None:
+    def remove_listener(
+        self, event: str, f: Callable[..., Union[Awaitable[None], None]]
+    ) -> None:
         """Removes the function ``f`` from ``event``."""
         self._impl_obj.remove_listener(event, self._wrap_handler(f))
 
@@ -93,4 +100,7 @@ async def __aexit__(
         exc_val: BaseException,
         traceback: TracebackType,
     ) -> None:
-        await self.close()  # type: ignore
+        await self.close()
+
+    async def close(self: Self) -> None:
+        ...
diff --git a/playwright/_impl/_sync_base.py b/playwright/_impl/_sync_base.py
@@ -36,18 +36,18 @@
 
 
 T = TypeVar("T")
-Self = TypeVar("Self")
+Self = TypeVar("Self", bound="SyncContextManager")
 
 
 class EventInfo(Generic[T]):
-    def __init__(self, sync_base: "SyncBase", future: asyncio.Future) -> None:
+    def __init__(self, sync_base: "SyncBase", future: "asyncio.Future[T]") -> None:
         self._sync_base = sync_base
         self._value: Optional[T] = None
-        self._exception = None
+        self._exception: Optional[Exception] = None
         self._future = future
         g_self = greenlet.getcurrent()
 
-        def done_callback(task: Any) -> None:
+        def done_callback(task: "asyncio.Future[T]") -> None:
             try:
                 self._value = mapping.from_maybe_impl(self._future.result())
             except Exception as e:
@@ -71,13 +71,18 @@ def is_done(self) -> bool:
 
 
 class EventContextManager(Generic[T]):
-    def __init__(self, sync_base: "SyncBase", future: asyncio.Future) -> None:
-        self._event: EventInfo = EventInfo(sync_base, future)
+    def __init__(self, sync_base: "SyncBase", future: "asyncio.Future[T]") -> None:
+        self._event = EventInfo[T](sync_base, future)
 
     def __enter__(self) -> EventInfo[T]:
         return self._event
 
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+    def __exit__(
+        self,
+        exc_type: Type[BaseException],
+        exc_val: BaseException,
+        exc_tb: TracebackType,
+    ) -> None:
         self._event.value
 
 
@@ -110,17 +115,17 @@ def _wrap_handler(self, handler: Any) -> Callable[..., None]:
             return mapping.wrap_handler(handler)
         return handler
 
-    def on(self, event: str, f: Any) -> None:
+    def on(self, event: str, f: Callable[..., None]) -> None:
         """Registers the function ``f`` to the event name ``event``."""
         self._impl_obj.on(event, self._wrap_handler(f))
 
-    def once(self, event: str, f: Any) -> None:
+    def once(self, event: str, f: Callable[..., None]) -> None:
         """The same as ``self.on``, except that the listener is automatically
         removed after being called.
         """
         self._impl_obj.once(event, self._wrap_handler(f))
 
-    def remove_listener(self, event: str, f: Any) -> None:
+    def remove_listener(self, event: str, f: Callable[..., None]) -> None:
         """Removes the function ``f`` from ``event``."""
         self._impl_obj.remove_listener(event, self._wrap_handler(f))
 
@@ -167,4 +172,7 @@ def __exit__(
         exc_val: BaseException,
         traceback: TracebackType,
     ) -> None:
-        self.close()  # type: ignore
+        self.close()
+
+    def close(self: Self) -> None:
+        ...
