[py] PEP 484 type hints for selenium.webdriver.remote.errorhandler

py/selenium/webdriver/remote/errorhandler.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from typing import Any, Dict, Mapping, Type, TypeVar
+
 from selenium.common.exceptions import (ElementClickInterceptedException,
                                         ElementNotInteractableException,
                                         ElementNotSelectableException,
@@ -45,6 +47,10 @@
                                         WebDriverException)
 
 
+_KT = TypeVar("_KT")
+_VT = TypeVar("_VT")
+
+
 class ErrorCode(object):
     """
     Error codes defined in the WebDriver wire protocol.
@@ -95,7 +101,7 @@ class ErrorHandler(object):
     Handles errors returned by the WebDriver server.
     """
 
-    def check_response(self, response):
+    def check_response(self, response: Dict[str, Any]) -> None:
         """
         Checks that a JSON response from the WebDriver does not have an error.
 
@@ -110,7 +116,7 @@ def check_response(self, response):
             return
         value = None
         message = response.get("message", "")
-        screen = response.get("screen", "")
+        screen: str = response.get("screen", "")
         stacktrace = None
         if isinstance(status, int):
             value_json = response.get('value', None)
@@ -132,6 +138,7 @@ def check_response(self, response):
                 except ValueError:
                     pass
 
+        exception_class: Type[WebDriverException]
         if status in ErrorCode.NO_SUCH_ELEMENT:
             exception_class = NoSuchElementException
         elif status in ErrorCode.NO_SUCH_FRAME:
@@ -201,7 +208,7 @@ def check_response(self, response):
         if message == "" and 'message' in value:
             message = value['message']
 
-        screen = None
+        screen = None  # type: ignore[assignment]
         if 'screen' in value:
             screen = value['screen']
 
@@ -232,8 +239,8 @@ def check_response(self, response):
                 alert_text = value['data'].get('text')
             elif 'alert' in value:
                 alert_text = value['alert'].get('text')
-            raise exception_class(message, screen, stacktrace, alert_text)
+            raise exception_class(message, screen, stacktrace, alert_text)  # type: ignore[call-arg]  # mypy is not smart enough here
         raise exception_class(message, screen, stacktrace)
 
-    def _value_or_default(self, obj, key, default):
+    def _value_or_default(self, obj: Mapping[_KT, _VT], key: _KT, default: _VT) -> _VT:
         return obj[key] if key in obj else default