speed up add get_account_ids

setup.py

@@ -5,15 +5,15 @@
 
 setuptools.setup(
     name="vanguard-api",
-    version="0.0.7",
+    version="0.0.8",
     author="MaxxRK",
     author_email="[email protected]",
     description="An unofficial API for Vanguard Invest",
     long_description=long_description,
     long_description_content_type="text/markdown",
     license="MIT",
     url="https://github.com/MaxxRK/vanguard-api",
-    download_url="https://github.com/MaxxRK/vanguard-api/archive/refs/tags/v0.0.7.tar.gz",
+    download_url="https://github.com/MaxxRK/vanguard-api/archive/refs/tags/v0.0.8.tar.gz",
     keywords=["VANGUARD", "API"],
     install_requires=["playwright", "playwright-stealth"],
     packages=["vanguard"],

vanguard/account.py

@@ -138,6 +138,29 @@ def _parse_rows(self, table_rows, account_id):
                     )
                     self.accounts_positions[account_id][type] = stocks
 
+    def get_account_ids(self):
+        """
+        Retrieves and sets the account numbers associated with the session.
+
+        This method navigates to the account holdings page, waits for the account numbers to load, and then retrieves the account numbers from the page.
+
+        Returns:
+            bool: True if the account numbers were successfully retrieved, False otherwise.
+        """
+        try:
+            self.session.go_url(holdings_page())
+            self.session.page.wait_for_selector(
+                '//span[contains(text(), "Expand all accounts")]', timeout=120000
+            ).click()
+            self.session.page.wait_for_selector("#overflow-override")
+            all_selectors = self.session.page.query_selector_all("#overflow-override")
+            for _, selector in enumerate(all_selectors):
+                account_id = self._get_account_id(selector)
+                self.account_numbers.append(account_id)
+            return True
+        except PlaywrightTimeoutError:
+            return False
+
     def get_holdings(self):
         """
         Retrieves and sets the holdings information of the account.
@@ -163,7 +186,6 @@ def get_holdings(self):
             all_selectors = self.session.page.query_selector_all("#overflow-override")
             for i, selector in enumerate(all_selectors):
                 account_id = self._get_account_id(selector)
-                self.account_numbers.append(account_id)
                 table_wrapper = selector.wait_for_selector(f"#self_managed_table_{i}")
                 table_entries = table_wrapper.query_selector_all("tbody")
                 for j, entry in enumerate(table_entries):

vanguard/session.py

@@ -145,14 +145,20 @@ def login(self, username, password, last_four):
         try:
             self.password = password
             self.go_url(landing_page())
-            try:
-                self.page.wait_for_selector(
-                    "#username-password-submit-btn-1", timeout=30000
-                )
-            except PlaywrightTimeoutError:
-                if self.page.url == landing_page():
-                    return False
-                raise Exception("Could not find submit button on login page.")
+            for _ in range(30):
+                try:
+                    if self.page.url == landing_page():
+                        self.page.wait_for_selector(
+                            "//h2[contains(text(), 'Accounts')]",
+                            timeout=1000,
+                        )
+                        return False
+                    self.page.wait_for_selector(
+                        "#username-password-submit-btn-1", timeout=1000
+                    )
+                    break
+                except PlaywrightTimeoutError:
+                    continue
             username_box = self.page.query_selector("#USER")
             username_box.type(username, delay=random.randint(50, 500))
             username_box.press("Tab")
@@ -162,13 +168,13 @@ def login(self, username, password, last_four):
             self.page.query_selector("#username-password-submit-btn-1").click()
             try:
                 self.page.wait_for_selector(
-                    "button.col-md:nth-child(2) > div:nth-child(1)", timeout=10000
+                    "button.col-md:nth-child(2) > div:nth-child(1)", timeout=5000
                 ).click()
             except PlaywrightTimeoutError:
                 pass
             try:
                 self.page.wait_for_selector(
-                    "xpath=//div[contains(text(), '***-***-')]", timeout=10000
+                    "xpath=//div[contains(text(), '***-***-')]", timeout=5000
                 )
                 otp_cards = self.page.query_selector_all(
                     "xpath=//div[contains(text(), '***-***-')]"
@@ -181,7 +187,7 @@ def login(self, username, password, last_four):
                 pass
             try:
                 self.page.wait_for_selector(
-                    "xpath=//div[contains(text(), 'Text')]", timeout=10000
+                    "xpath=//div[contains(text(), 'Text')]", timeout=5000
                 ).click()
                 return True
             except PlaywrightTimeoutError:
@@ -208,17 +214,17 @@ def login_two(self, code):
         """
         try:
             code = str(code)
-            self.page.wait_for_selector("#CODE", timeout=10000).fill(code)
+            self.page.wait_for_selector("#CODE", timeout=5000).fill(code)
             self.page.query_selector(
                 "c11n-radio.c11n-radio:nth-child(2) > label:nth-child(2)"
             ).click()
             self.page.wait_for_selector(
-                "#security-code-submit-btn", timeout=10000
+                "#security-code-submit-btn", timeout=5000
             ).click()
             sleep(5)
             try:
                 self.page.wait_for_url(
-                    landing_page(), wait_until="domcontentloaded", timeout=10000
+                    landing_page(), wait_until="domcontentloaded", timeout=5000
                 )
                 if self.title is not None:
                     self.save_storage_state()