From fd834b7768497fc0a6561b0afcf79fa533f9ba14 Mon Sep 17 00:00:00 2001 From: pinterior Date: Thu, 2 Nov 2023 03:26:57 +0900 Subject: [PATCH] more strict typing on s.w.support.relative_locator --- py/selenium/webdriver/common/by.py | 5 ++ .../webdriver/support/relative_locator.py | 61 ++++++++++++++++--- 2 files changed, 58 insertions(+), 8 deletions(-) diff --git a/py/selenium/webdriver/common/by.py b/py/selenium/webdriver/common/by.py index f24a113ddb1bec..56a3f96d6fbb84 100644 --- a/py/selenium/webdriver/common/by.py +++ b/py/selenium/webdriver/common/by.py @@ -16,6 +16,8 @@ # under the License. """The By implementation.""" +from typing import Literal + class By: """Set of supported locator strategies.""" @@ -28,3 +30,6 @@ class By: TAG_NAME = "tag name" CLASS_NAME = "class name" CSS_SELECTOR = "css selector" + + +ByType = Literal["id", "xpath", "link text", "partial link text", "name", "tag name", "class name", "css selector"] diff --git a/py/selenium/webdriver/support/relative_locator.py b/py/selenium/webdriver/support/relative_locator.py index 5dab29c3410fd1..f5976e6debf492 100644 --- a/py/selenium/webdriver/support/relative_locator.py +++ b/py/selenium/webdriver/support/relative_locator.py @@ -16,11 +16,14 @@ # under the License. from typing import Dict from typing import List +from typing import NoReturn from typing import Optional from typing import Union +from typing import overload from selenium.common.exceptions import WebDriverException from selenium.webdriver.common.by import By +from selenium.webdriver.common.by import ByType from selenium.webdriver.remote.webelement import WebElement @@ -37,10 +40,10 @@ def with_tag_name(tag_name: str) -> "RelativeBy": """ if not tag_name: raise WebDriverException("tag_name can not be null") - return RelativeBy({"css selector": tag_name}) + return RelativeBy({By.CSS_SELECTOR: tag_name}) -def locate_with(by: By, using: str) -> "RelativeBy": +def locate_with(by: ByType, using: str) -> "RelativeBy": """Start searching for relative objects your search criteria with By. :Args: @@ -70,7 +73,9 @@ class RelativeBy: assert "mid" in ids """ - def __init__(self, root: Optional[Dict[Union[By, str], str]] = None, filters: Optional[List] = None): + LocatorType = Dict[ByType, str] + + def __init__(self, root: Optional[Dict[ByType, str]] = None, filters: Optional[List] = None): """Creates a new RelativeBy object. It is preferred if you use the `locate_with` method as this signature could change. @@ -82,7 +87,15 @@ def __init__(self, root: Optional[Dict[Union[By, str], str]] = None, filters: Op self.root = root self.filters = filters or [] - def above(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy": + @overload + def above(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": + ... + + @overload + def above(self, element_or_locator: None = None) -> "NoReturn": + ... + + def above(self, element_or_locator: Union[WebElement, LocatorType, None] = None) -> "RelativeBy": """Add a filter to look for elements above. :Args: @@ -94,7 +107,15 @@ def above(self, element_or_locator: Union[WebElement, Dict] = None) -> "Relative self.filters.append({"kind": "above", "args": [element_or_locator]}) return self - def below(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy": + @overload + def below(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": + ... + + @overload + def below(self, element_or_locator: None = None) -> "NoReturn": + ... + + def below(self, element_or_locator: Union[WebElement, Dict, None] = None) -> "RelativeBy": """Add a filter to look for elements below. :Args: @@ -106,7 +127,15 @@ def below(self, element_or_locator: Union[WebElement, Dict] = None) -> "Relative self.filters.append({"kind": "below", "args": [element_or_locator]}) return self - def to_left_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy": + @overload + def to_left_of(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": + ... + + @overload + def to_left_of(self, element_or_locator: None = None) -> "NoReturn": + ... + + def to_left_of(self, element_or_locator: Union[WebElement, Dict, None] = None) -> "RelativeBy": """Add a filter to look for elements to the left of. :Args: @@ -118,7 +147,15 @@ def to_left_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "Rel self.filters.append({"kind": "left", "args": [element_or_locator]}) return self - def to_right_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy": + @overload + def to_right_of(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": + ... + + @overload + def to_right_of(self, element_or_locator: None = None) -> "NoReturn": + ... + + def to_right_of(self, element_or_locator: Union[WebElement, Dict, None] = None) -> "RelativeBy": """Add a filter to look for elements right of. :Args: @@ -130,7 +167,15 @@ def to_right_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "Re self.filters.append({"kind": "right", "args": [element_or_locator]}) return self - def near(self, element_or_locator: Union[WebElement, Dict, None] = None, distance: int = 50) -> "RelativeBy": + @overload + def near(self, element_or_locator: Union[WebElement, LocatorType], distance: int = 50) -> "RelativeBy": + ... + + @overload + def near(self, element_or_locator: None = None, distance: int = 50) -> "NoReturn": + ... + + def near(self, element_or_locator: Union[WebElement, LocatorType, None] = None, distance: int = 50) -> "RelativeBy": """Add a filter to look for elements near. :Args: