Make convenience updates to "driver.find_element()"

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -107,7 +107,10 @@ def make_driver_executable_if_not(driver_path):
 
 def extend_driver(driver):
     # Extend the driver with new methods
+    driver.default_find_element = driver.find_element
     DM = sb_driver.DriverMethods(driver)
+    driver.find_element = DM.find_element
+    driver.locator = DM.locator
     page = types.SimpleNamespace()
     page.open = DM.open_url
     page.click = DM.click

seleniumbase/core/sb_driver.py

@@ -1,12 +1,34 @@
 """Add new methods to extend the driver"""
 from seleniumbase.fixtures import js_utils
 from seleniumbase.fixtures import page_actions
+from seleniumbase.fixtures import page_utils
 
 
 class DriverMethods():
     def __init__(self, driver):
         self.driver = driver
 
+    def find_element(self, by, value=None):
+        if not value:
+            value = by
+            by = "css selector"
+        else:
+            value, by = page_utils.swap_selector_and_by_if_reversed(value, by)
+        return self.driver.default_find_element(by=by, value=value)
+
+    def locator(self, selector, by=None):
+        if not by:
+            by = "css selector"
+        else:
+            selector, by = page_utils.swap_selector_and_by_if_reversed(
+                selector, by
+            )
+        try:
+            return self.driver.default_find_element(by=by, value=selector)
+        except Exception:
+            pass
+        raise Exception('No such Element: {%s} (by="%s")!' % (selector, by))
+
     def open_url(self, *args, **kwargs):
         page_actions.open_url(self.driver, *args, **kwargs)