From 7c9246887410636220dc134b4a47f5291aaf962b Mon Sep 17 00:00:00 2001
From: Simon K <jackofspaces@gmail.com>
Date: Thu, 9 Jun 2022 13:56:17 +0100
Subject: [PATCH] [py] Additional types and tweaks for remote webdriver
 (#10631)

* [py] Tidy some remote webdriver code

* [py]: revert imports in remote webdriver; more types and tweaks

* [py]: Fix `flake8` violations in remote webdriver
---
 py/selenium/webdriver/remote/file_detector.py |  2 +-
 py/selenium/webdriver/remote/webdriver.py     | 47 +++++++------------
 2 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/py/selenium/webdriver/remote/file_detector.py b/py/selenium/webdriver/remote/file_detector.py
index 8499f95bf99a7..a0e15f2bdb76a 100644
--- a/py/selenium/webdriver/remote/file_detector.py
+++ b/py/selenium/webdriver/remote/file_detector.py
@@ -30,7 +30,7 @@ class FileDetector(metaclass=ABCMeta):
 
     @abstractmethod
     def is_local_file(self, *keys: AnyKey) -> Optional[str]:
-        return None
+        raise NotImplementedError
 
 
 class UselessFileDetector(FileDetector):
diff --git a/py/selenium/webdriver/remote/webdriver.py b/py/selenium/webdriver/remote/webdriver.py
index 6dac03419f777..8572761d46683 100644
--- a/py/selenium/webdriver/remote/webdriver.py
+++ b/py/selenium/webdriver/remote/webdriver.py
@@ -16,8 +16,10 @@
 # under the License.
 
 """The WebDriver implementation."""
-
+import contextlib
 import copy
+import types
+import typing
 from importlib import import_module
 
 import pkgutil
@@ -281,7 +283,10 @@ def __repr__(self):
     def __enter__(self):
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self,
+                 exc_type: typing.Optional[typing.Type[BaseException]],
+                 exc: typing.Optional[BaseException],
+                 traceback: typing.Optional[types.TracebackType]):
         self.quit()
 
     @contextmanager
@@ -314,7 +319,7 @@ def file_detector_context(self, file_detector_class, *args, **kwargs):
                 self.file_detector = last_detector
 
     @property
-    def mobile(self):
+    def mobile(self) -> Mobile:
         return self._mobile
 
     @property
@@ -450,30 +455,17 @@ def title(self) -> str:
 
                 title = driver.title
         """
-        resp = self.execute(Command.GET_TITLE)
-        return resp['value'] if resp['value'] else ""
+        return self.execute(Command.GET_TITLE).get("value", "")
 
     def pin_script(self, script, script_key=None) -> ScriptKey:
-        """
-
-        """
-        if not script_key:
-            _script_key = ScriptKey()
-        else:
-            _script_key = ScriptKey(script_key)
+        _script_key = ScriptKey(script_key)
         self.pinned_scripts[_script_key.id] = script
         return _script_key
 
     def unpin(self, script_key) -> None:
-        """
-
-        """
         self.pinned_scripts.pop(script_key.id)
 
     def get_pinned_scripts(self) -> List[str]:
-        """
-
-        """
         return list(self.pinned_scripts.keys())
 
     def execute_script(self, script, *args):
@@ -602,9 +594,8 @@ def maximize_window(self) -> None:
         """
         Maximizes the current window that webdriver is using
         """
-        params = None
         command = Command.W3C_MAXIMIZE_WINDOW
-        self.execute(command, params)
+        self.execute(command, None)
 
     def fullscreen_window(self) -> None:
         """
@@ -695,7 +686,7 @@ def get_cookies(self) -> List[dict]:
         """
         return self.execute(Command.GET_ALL_COOKIES)['value']
 
-    def get_cookie(self, name) -> dict:
+    def get_cookie(self, name) -> typing.Optional[typing.Dict]:
         """
         Get a single cookie by name. Returns the cookie if found, None if not.
 
@@ -704,10 +695,8 @@ def get_cookie(self, name) -> dict:
 
                 driver.get_cookie('my_cookie')
         """
-        try:
-            return self.execute(Command.GET_COOKIE, {'name': name})['value']
-        except NoSuchCookieException:
-            return None
+        with contextlib.suppress(NoSuchCookieException):
+            return self.execute(Command.GET_COOKIE, {"name": name})['value']
 
     def delete_cookie(self, name) -> None:
         """
@@ -986,7 +975,7 @@ def get_screenshot_as_base64(self) -> str:
         """
         return self.execute(Command.SCREENSHOT)['value']
 
-    def set_window_size(self, width, height, windowHandle='current') -> dict:
+    def set_window_size(self, width, height, windowHandle: str = 'current') -> None:
         """
         Sets the width and height of the current window. (window.resizeTo)
 
@@ -1003,7 +992,7 @@ def set_window_size(self, width, height, windowHandle='current') -> dict:
             warnings.warn("Only 'current' window is supported for W3C compatible browsers.")
         self.set_window_rect(width=int(width), height=int(height))
 
-    def get_window_size(self, windowHandle='current') -> dict:
+    def get_window_size(self, windowHandle: str = 'current') -> dict:
         """
         Gets the width and height of the current window.
 
@@ -1022,7 +1011,7 @@ def get_window_size(self, windowHandle='current') -> dict:
 
         return {k: size[k] for k in ('width', 'height')}
 
-    def set_window_position(self, x, y, windowHandle='current') -> dict:
+    def set_window_position(self, x, y, windowHandle: str = 'current') -> dict:
         """
         Sets the x,y position of the current window. (window.moveTo)
 
@@ -1090,7 +1079,7 @@ def set_window_rect(self, x=None, y=None, width=None, height=None) -> dict:
                                                       "height": height})['value']
 
     @property
-    def file_detector(self):
+    def file_detector(self) -> FileDetector:
         return self._file_detector
 
     @file_detector.setter