WIP: Added more type hints - needs review

Author: DeflateAwning

py/selenium/webdriver/common/actions/action_builder.py

@@ -103,7 +103,7 @@ def add_key_input(self, name: str) -> KeyInput:
         self._add_input(new_input)
         return new_input
 
-    def add_pointer_input(self, kind: str, name: str) -> PointerInput:
+    def add_pointer_input(self, kind: interaction.POINTER_KINDS_LITERAL, name: str) -> PointerInput:
         """Add a new pointer input device to the action builder.
 
         Parameters:

py/selenium/webdriver/common/actions/input_device.py

@@ -19,6 +19,7 @@
 from typing import Any
 from typing import List
 from typing import Optional
+from typing import Union
 
 
 class InputDevice:
@@ -35,5 +36,5 @@ def add_action(self, action: Any) -> None:
     def clear_actions(self) -> None:
         self.actions = []
 
-    def create_pause(self, duration: float = 0) -> None:
+    def create_pause(self, pause_duration: Union[float, int] = 0) -> None:
         pass

py/selenium/webdriver/common/actions/interaction.py

@@ -15,30 +15,39 @@
 # specific language governing permissions and limitations
 # under the License.
 from typing import Dict
+from typing import Literal
 from typing import Union
 
+from .key_input import KeyInput
+from .pointer_input import PointerInput
+from .wheel_input import WheelInput
+
 KEY = "key"
 POINTER = "pointer"
 NONE = "none"
 WHEEL = "wheel"
 SOURCE_TYPES = {KEY, POINTER, NONE}
+SOURCE_TYPES_LITERAL = Literal['key', 'pointer', 'none', 'wheel']
 
 POINTER_MOUSE = "mouse"
 POINTER_TOUCH = "touch"
 POINTER_PEN = "pen"
-
 POINTER_KINDS = {POINTER_MOUSE, POINTER_TOUCH, POINTER_PEN}
+POINTER_KINDS_LITERAL = Literal['mouse', 'touch', 'pen']
 
 
 class Interaction:
     PAUSE = "pause"
 
-    def __init__(self, source: str) -> None:
+    def __init__(self, source: Union[KeyInput, PointerInput, WheelInput]) -> None:
         self.source = source
 
 
 class Pause(Interaction):
-    def __init__(self, source, duration: float = 0) -> None:
+    def __init__(self, source: Union[KeyInput, PointerInput, WheelInput], duration: Union[int, float] = 0) -> None:
+        """
+        :Args:
+         - duration: duration of the pause in seconds"""
         super().__init__(source)
         self.duration = duration
 

py/selenium/webdriver/common/actions/key_input.py

@@ -14,36 +14,42 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+from typing import Any
+from typing import Literal
+from typing import Union
+
 from . import interaction
 from .input_device import InputDevice
 from .interaction import Interaction
 from .interaction import Pause
+from .interaction import SOURCE_TYPES_LITERAL
 
 
 class KeyInput(InputDevice):
     def __init__(self, name: str) -> None:
         super().__init__()
         self.name = name
-        self.type = interaction.KEY
+        self.type: SOURCE_TYPES_LITERAL = interaction.KEY
 
-    def encode(self) -> dict:
+    def encode(self) -> dict[str, Any]:
         return {"type": self.type, "id": self.name, "actions": [acts.encode() for acts in self.actions]}
 
-    def create_key_down(self, key) -> None:
+    def create_key_down(self, key: str) -> None:
         self.add_action(TypingInteraction(self, "keyDown", key))
 
-    def create_key_up(self, key) -> None:
+    def create_key_up(self, key: str) -> None:
         self.add_action(TypingInteraction(self, "keyUp", key))
 
-    def create_pause(self, pause_duration: float = 0) -> None:
+    def create_pause(self, pause_duration: Union[float, int] = 0) -> None:
         self.add_action(Pause(self, pause_duration))
 
 
 class TypingInteraction(Interaction):
-    def __init__(self, source, type_, key) -> None:
+    def __init__(self, source: str, type_: Literal["keyUp", "keyDown"], key: str) -> None:
         super().__init__(source)
         self.type = type_
         self.key = key
 
-    def encode(self) -> dict:
+    def encode(self) -> dict[str, str]:
         return {"type": self.type, "value": self.key}

py/selenium/webdriver/common/actions/pointer_input.py

@@ -26,51 +26,52 @@
 from .input_device import InputDevice
 from .interaction import POINTER
 from .interaction import POINTER_KINDS
+from .interaction import POINTER_KINDS_LITERAL, SOURCE_TYPES_LITERAL
 
 
 class PointerInput(InputDevice):
     DEFAULT_MOVE_DURATION = 250
 
-    def __init__(self, kind, name):
+    def __init__(self, kind: POINTER_KINDS_LITERAL, name: str) -> None:
         super().__init__()
         if kind not in POINTER_KINDS:
             raise InvalidArgumentException(f"Invalid PointerInput kind '{kind}'")
-        self.type = POINTER
-        self.kind = kind
+        self.type: SOURCE_TYPES_LITERAL = POINTER
+        self.kind: POINTER_KINDS_LITERAL = kind
         self.name = name
 
     def create_pointer_move(
         self,
-        duration=DEFAULT_MOVE_DURATION,
-        x: float = 0,
-        y: float = 0,
-        origin: Optional[WebElement] = None,
+        duration: Union[int, float] = DEFAULT_MOVE_DURATION,
+        x: Union[int, float] = 0,
+        y: Union[int, float] = 0,
+        origin: Union[WebElement, str, None] = None,
         **kwargs,
-    ):
+    ) -> None:
         action = {"type": "pointerMove", "duration": duration, "x": x, "y": y, **kwargs}
         if isinstance(origin, WebElement):
             action["origin"] = {"element-6066-11e4-a52e-4f735466cecf": origin.id}
         elif origin is not None:
             action["origin"] = origin
         self.add_action(self._convert_keys(action))
 
-    def create_pointer_down(self, **kwargs):
+    def create_pointer_down(self, **kwargs) -> None:
         data = {"type": "pointerDown", "duration": 0, **kwargs}
         self.add_action(self._convert_keys(data))
 
-    def create_pointer_up(self, button):
+    def create_pointer_up(self, button) -> None:
         self.add_action({"type": "pointerUp", "duration": 0, "button": button})
 
-    def create_pointer_cancel(self):
+    def create_pointer_cancel(self) -> None:
         self.add_action({"type": "pointerCancel"})
 
     def create_pause(self, pause_duration: Union[int, float] = 0) -> None:
         self.add_action({"type": "pause", "duration": int(pause_duration * 1000)})
 
-    def encode(self):
+    def encode(self) -> Dict[str, Any]:
         return {"type": self.type, "parameters": {"pointerType": self.kind}, "id": self.name, "actions": self.actions}
 
-    def _convert_keys(self, actions: Dict[str, Any]):
+    def _convert_keys(self, actions: Dict[str, Any]) -> Dict[str, Any]:
         out = {}
         for k, v in actions.items():
             if v is None:

py/selenium/webdriver/common/actions/wheel_actions.py

@@ -16,21 +16,34 @@
 # under the License.
 
 from typing import Optional
+from typing import Union
 
 from .interaction import Interaction
 from .wheel_input import WheelInput
 
 
 class WheelActions(Interaction):
-    def __init__(self, source: Optional[WheelInput] = None):
-        if not source:
+    def __init__(self, source: Optional[WheelInput] = None) -> None:
+        if source is None:
             source = WheelInput("wheel")
         super().__init__(source)
 
-    def pause(self, duration: float = 0):
+    def pause(self, duration: Union[float, int] = 0) -> "WheelActions":
         self.source.create_pause(duration)
         return self
 
-    def scroll(self, x=0, y=0, delta_x=0, delta_y=0, duration=0, origin="viewport"):
+    def scroll(
+        self,
+        x: int = 0,
+        y: int = 0,
+        delta_x: int = 0,
+        delta_y: int = 0,
+        duration: Union[float, int] = 0,
+        origin: str = "viewport",
+    ) -> "WheelActions":
+        """
+        :Args:
+         - duration: The duration of the scroll, in seconds.
+        """
         self.source.create_scroll(x, y, delta_x, delta_y, duration, origin)
         return self

py/selenium/webdriver/common/actions/wheel_input.py

@@ -14,6 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+from typing import Any
+from typing import Dict
 from typing import Union
 
 from selenium.webdriver.remote.webelement import WebElement
@@ -29,11 +32,11 @@ def __init__(self, origin: Union[str, WebElement], x_offset: int, y_offset: int)
         self._y_offset = y_offset
 
     @classmethod
-    def from_element(cls, element: WebElement, x_offset: int = 0, y_offset: int = 0):
+    def from_element(cls, element: WebElement, x_offset: int = 0, y_offset: int = 0) -> "ScrollOrigin":
         return cls(element, x_offset, y_offset)
 
     @classmethod
-    def from_viewport(cls, x_offset: int = 0, y_offset: int = 0):
+    def from_viewport(cls, x_offset: int = 0, y_offset: int = 0) -> "ScrollOrigin":
         return cls("viewport", x_offset, y_offset)
 
     @property
@@ -50,15 +53,21 @@ def y_offset(self) -> int:
 
 
 class WheelInput(InputDevice):
-    def __init__(self, name) -> None:
+    def __init__(self, name: str) -> None:
         super().__init__(name=name)
         self.name = name
-        self.type = interaction.WHEEL
+        self.type: interaction.SOURCE_TYPES_LITERAL = interaction.WHEEL
 
-    def encode(self) -> dict:
+    def encode(self) -> dict[str, Any]:
         return {"type": self.type, "id": self.name, "actions": self.actions}
 
-    def create_scroll(self, x: int, y: int, delta_x: int, delta_y: int, duration: int, origin) -> None:
+    def create_scroll(
+        self, x: int, y: int, delta_x: int, delta_y: int, duration: Union[float, int], origin: Union[str, WebElement, Dict[str, str]]
+    ) -> None:
+        """
+        :Args:
+         - duration: The duration to pause for after the scroll, in seconds.
+        """
         if isinstance(origin, WebElement):
             origin = {"element-6066-11e4-a52e-4f735466cecf": origin.id}
         self.add_action(
@@ -74,4 +83,8 @@ def create_scroll(self, x: int, y: int, delta_x: int, delta_y: int, duration: in
         )
 
     def create_pause(self, pause_duration: Union[int, float] = 0) -> None:
+        """
+        :Args:
+         - pause_duration: duration of the pause in seconds
+        """
         self.add_action({"type": "pause", "duration": int(pause_duration * 1000)})