custom overload decorator to support overloads in the api generate scripts

Author: DetachHead

playwright/_impl/_element_handle.py

@@ -33,6 +33,7 @@
     parse_result,
     serialize_argument,
 )
+from playwright._impl._overload import api_overload
 
 if sys.version_info >= (3, 8):  # pragma: no cover
     from typing import Literal
@@ -309,7 +310,8 @@ async def wait_for_element_state(
     ) -> None:
         await self._channel.send("waitForElementState", locals_to_params(locals()))
 
-    async def _overload1_wait_for_selector(
+    @api_overload
+    async def wait_for_selector(
         self,
         selector: str,
         *,
@@ -318,7 +320,8 @@ async def _overload1_wait_for_selector(
     ) -> "ElementHandle":
         ...
 
-    async def _overload2_wait_for_selector(
+    @api_overload  # type: ignore[no-redef]
+    async def wait_for_selector(
         self,
         selector: str,
         *,
@@ -327,7 +330,7 @@ async def _overload2_wait_for_selector(
     ) -> None:
         ...
 
-    async def wait_for_selector(
+    async def wait_for_selector(  # type: ignore[no-redef]
         self,
         selector: str,
         *,

playwright/_impl/_frame.py

@@ -48,6 +48,7 @@
 )
 from playwright._impl._locator import Locator
 from playwright._impl._network import Response
+from playwright._impl._overload import api_overload
 from playwright._impl._wait_helper import WaitHelper
 
 if sys.version_info >= (3, 8):  # pragma: no cover
@@ -270,7 +271,8 @@ async def query_selector_all(self, selector: str) -> List[ElementHandle]:
             )
         )
 
-    async def _overload1_wait_for_selector(
+    @api_overload
+    async def wait_for_selector(
         self,
         selector: str,
         *,
@@ -280,7 +282,8 @@ async def _overload1_wait_for_selector(
     ) -> ElementHandle:
         ...
 
-    async def _overload2_wait_for_selector(
+    @api_overload  # type: ignore[no-redef]
+    async def wait_for_selector(
         self,
         selector: str,
         *,
@@ -290,7 +293,7 @@ async def _overload2_wait_for_selector(
     ) -> None:
         ...
 
-    async def wait_for_selector(
+    async def wait_for_selector(  # type: ignore[no-redef]
         self,
         selector: str,
         *,

playwright/_impl/_overload.py

@@ -0,0 +1,16 @@
+from collections import defaultdict
+import inspect
+from typing import Callable
+
+
+def api_overload(fn: Callable) -> Callable:
+    """
+    Creates an overload that's visible to the api generate scripts. use this decorator instead of `typing.overload` when exposing overloads from `_impl`.
+    You will need to suppress mypy errors using a `# type: ignore[no-redef]` comment
+    """
+    dictionary = inspect.getmodule(fn).__dict__
+    overloads_key = "__overloads__"
+    if dictionary.get(overloads_key) is None:
+        dictionary[overloads_key] = defaultdict(list)
+    dictionary[overloads_key][fn.__name__].append(fn)
+    return fn

playwright/_impl/_page.py

@@ -71,6 +71,7 @@
     serialize_argument,
 )
 from playwright._impl._network import Request, Response, Route, serialize_headers
+from playwright._impl._overload import api_overload
 from playwright._impl._video import Video
 from playwright._impl._wait_helper import WaitHelper
 
@@ -325,7 +326,8 @@ async def query_selector(
     async def query_selector_all(self, selector: str) -> List[ElementHandle]:
         return await self._main_frame.query_selector_all(selector)
 
-    async def _overload1_wait_for_selector(
+    @api_overload
+    async def wait_for_selector(
         self,
         selector: str,
         *,
@@ -335,7 +337,8 @@ async def _overload1_wait_for_selector(
     ) -> ElementHandle:
         ...
 
-    async def _overload2_wait_for_selector(
+    @api_overload  # type: ignore[no-redef]
+    async def wait_for_selector(
         self,
         selector: str,
         *,
@@ -345,7 +348,7 @@ async def _overload2_wait_for_selector(
     ) -> None:
         ...
 
-    async def wait_for_selector(
+    async def wait_for_selector(  # type: ignore[no-redef]
         self,
         selector: str,
         *,

scripts/generate_api.py

@@ -19,7 +19,6 @@
     Callable,
     List,
     Match,
-    Optional,
     Union,
     cast,
     get_args,
@@ -93,34 +92,6 @@ def process_type(value: Any, param: bool = False) -> str:
 ]
 
 
-class Overload:
-    def __init__(self, clazz: object, name: str):
-        matches = re.match(r"_overload(\d+)_([a-zA-Z\d_]+)", name)
-        if matches is None:
-            raise ValueError(f"failed to get function name for overload: {name}")
-        self.number = int(matches[1])
-        self.name = matches[2]
-        self.clazz = clazz
-
-    def assert_has_implementation(self) -> None:
-        if self.implementation() is None:
-            raise Exception(f"implementation for overload '{self.name}' not found")
-
-    def implementation(self) -> Optional[Callable]:
-        for name, method in self.clazz.__dict__.items():
-            if name == self.name:
-                return method
-        return None
-
-
-def is_overload(name: str) -> bool:
-    try:
-        Overload(None, name)
-        return True
-    except TypeError:
-        return False
-
-
 def is_positional_exception(key: str) -> bool:
     for pattern in positional_exceptions:
         if re.match(pattern, key):
@@ -234,11 +205,15 @@ def arguments(func: FunctionType, indent: int) -> str:
     return split.join(tokens)
 
 
-def return_type(func: FunctionType) -> str:
+def return_type(func: Callable) -> str:
     value = get_type_hints(func, globals())["return"]
     return process_type(value)
 
 
+def return_type_value(func: Callable) -> str:
+    return re.sub(r"\"([^\"]+)Impl\"", r"\1", return_type(func))
+
+
 def short_name(t: Any) -> str:
     match = cast(
         Match[str], re.compile(r"playwright\._impl\.[^.]+\.([^']+)").search(str(t))

scripts/generate_async_api.py

@@ -14,21 +14,19 @@
 # limitations under the License.
 
 import inspect
-import re
 from types import FunctionType
-from typing import Any, Optional, get_type_hints
+from typing import Any, Callable, Dict, List, get_type_hints
 
 from playwright._impl._helper import to_snake_case
 from scripts.documentation_provider import DocumentationProvider
 from scripts.generate_api import (
-    Overload,
     all_types,
     api_globals,
     arguments,
     header,
-    is_overload,
     process_type,
     return_type,
+    return_type_value,
     return_value,
     short_name,
     signature,
@@ -37,6 +35,12 @@
 documentation_provider = DocumentationProvider(True)
 
 
+def async_return_type_value(func: Callable) -> str:
+    return return_type_value(func).replace(
+        "EventContextManager", "AsyncEventContextManager"
+    )
+
+
 def generate(t: Any) -> None:
     print("")
     class_name = short_name(t)
@@ -59,6 +63,9 @@ def generate(t: Any) -> None:
         [prefix, suffix] = return_value(type)
         prefix = "        return " + prefix + f"self._impl_obj.{name}"
         print(f"{prefix}{suffix}")
+    all_overloads: Dict[str, List[Callable]] = (
+        inspect.getmodule(t).__dict__.get("__overloads__") or {}
+    )
     for [name, value] in t.__dict__.items():
         if name.startswith("_"):
             continue
@@ -78,55 +85,48 @@ def generate(t: Any) -> None:
             prefix = "        return " + prefix + f"self._impl_obj.{name}"
             print(f"{prefix}{arguments(value, len(prefix))}{suffix}")
     for [name, value] in t.__dict__.items():
-        overload: Optional[Overload] = None
-        if is_overload(name):
-            overload = Overload(t, name)
-            overload.assert_has_implementation()
-            name = overload.name
+        overloads = all_overloads.get(name) or []
         if (
             not name.startswith("_")
             and isinstance(value, FunctionType)
             and "remove_listener" != name
         ):
             is_async = inspect.iscoroutinefunction(value)
-            return_type_value = return_type(value)
-            return_type_value = re.sub(r"\"([^\"]+)Impl\"", r"\1", return_type_value)
-            return_type_value = return_type_value.replace(
-                "EventContextManager", "AsyncEventContextManager"
-            )
-            print("")
             async_prefix = "async " if is_async else ""
-            if overload is not None:
+            indent = len(name) + 9
+            for overload in overloads:
                 print("    @typing.overload")
+                print(
+                    f"    {async_prefix}def {name}({signature(overload, indent, True)}) -> {async_return_type_value(overload)}:"
+                )
+                print("        pass")
+            print("")
             print(
-                f"    {async_prefix}def {name}({signature(value, len(name) + 9, overload is not None)}) -> {return_type_value}:"
+                f"    {async_prefix}def {name}({signature(value, indent)}) -> {async_return_type_value(value)}:"
             )
-            if overload is None:
-                documentation_provider.print_entry(
-                    class_name, name, get_type_hints(value, api_globals)
+            documentation_provider.print_entry(
+                class_name, name, get_type_hints(value, api_globals)
+            )
+            if "expect_" in name:
+                print("")
+                print(
+                    f"        return AsyncEventContextManager(self._impl_obj.{name}({arguments(value, 12)}).future)"
                 )
-                if "expect_" in name:
-                    print("")
-                    print(
-                        f"        return AsyncEventContextManager(self._impl_obj.{name}({arguments(value, 12)}).future)"
-                    )
-                else:
-                    [prefix, suffix] = return_value(
-                        get_type_hints(value, api_globals)["return"]
+            else:
+                [prefix, suffix] = return_value(
+                    get_type_hints(value, api_globals)["return"]
+                )
+                if is_async:
+                    prefix += (
+                        f'await self._async("{to_snake_case(class_name)}.{name}", '
                     )
-                    if is_async:
-                        prefix += (
-                            f'await self._async("{to_snake_case(class_name)}.{name}", '
-                        )
-                        suffix += ")"
-                    prefix = prefix + f"self._impl_obj.{name}("
-                    suffix = ")" + suffix
-                    print(
-                        f"""
+                    suffix += ")"
+                prefix = prefix + f"self._impl_obj.{name}("
+                suffix = ")" + suffix
+                print(
+                    f"""
         return {prefix}{arguments(value, len(prefix))}{suffix}"""
-                    )
-            else:
-                print("        pass")
+                )
     if class_name == "Playwright":
         print(
             """