Optimize the "find_visible_elements(selector)" method

Author: mdmintz

seleniumbase/fixtures/base_case.py

@@ -2162,9 +2162,9 @@ def find_visible_elements(self, selector, by="css selector", limit=0):
         selector, by = self.__recalculate_selector(selector, by)
         self.wait_for_ready_state_complete()
         time.sleep(0.05)
-        v_elems = page_actions.find_visible_elements(self.driver, selector, by)
-        if limit and limit > 0 and len(v_elems) > limit:
-            v_elems = v_elems[:limit]
+        v_elems = page_actions.find_visible_elements(
+            self.driver, selector, by, limit
+        )
         return v_elems
 
     def click_visible_elements(

seleniumbase/fixtures/page_actions.py

@@ -1251,26 +1251,40 @@ def wait_for_attribute_not_present(
     timeout_exception(Exception, message)
 
 
-def find_visible_elements(driver, selector, by="css selector"):
+def find_visible_elements(driver, selector, by="css selector", limit=0):
     """
     Finds all WebElements that match a selector and are visible.
-    Similar to webdriver.find_elements.
+    Similar to webdriver.find_elements().
+    If "limit" is set and > 0, will only return that many elements.
     @Params
     driver - the webdriver object (required)
     selector - the locator for identifying the page element (required)
     by - the type of selector being used (Default: "css selector")
+    limit - the maximum number of elements to return if > 0.
     """
     elements = driver.find_elements(by=by, value=selector)
+    if limit and limit > 0 and len(elements) > limit:
+        elements = elements[:limit]
     try:
         v_elems = [element for element in elements if element.is_displayed()]
         return v_elems
     except (StaleElementReferenceException, ElementNotInteractableException):
         time.sleep(0.1)
         elements = driver.find_elements(by=by, value=selector)
+        extra_elements = []
+        if limit and limit > 0 and len(elements) > limit:
+            elements = elements[:limit]
+            extra_elements = elements[limit:]
         v_elems = []
         for element in elements:
             if element.is_displayed():
                 v_elems.append(element)
+        if extra_elements and limit and len(v_elems) < limit:
+            for element in extra_elements:
+                if element.is_displayed():
+                    v_elems.append(element)
+                    if len(v_elems) >= limit:
+                        break
         return v_elems