Async for issue #11

.github/workflows/test.yml

@@ -0,0 +1,35 @@
+name: Test async
+
+jobs:
+  scheduled:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        name: Checkout repo
+      - name: Set up Python
+        uses: actions/setup-python@v1
+        with:
+          python-version: 3.11
+      - uses: actions/cache@v2
+        name: Configure pip caching
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-
+      - name: Install Python dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install playwright
+          python -m pip install playwright-stealth
+          python -m pip install pyotp
+          python -m pip install python-vipaccess
+          python -m pip install python-dotenv
+          python -m playwright install firefox
+      - name: run test
+        env:
+          PYTHONPATH: ${{ github.workspace }}
+          SCHWAB_USERNAME: ${{ secrets.SCHWAB_USERNAME }}
+          SCHWAB_PASSWORD: ${{ secrets.SCHWAB_PASSWORD }}
+          SCHWAB_TOTP: ${{ secrets.SCHWAB_TOTP }}
+        run: python example/example.py

example/example.py

@@ -61,4 +61,4 @@
 
 orders = api.orders_v2()
 
-pprint.pprint(orders)
+pprint.pprint(orders)

schwab_api/authentication.py

@@ -3,9 +3,10 @@
 import re
 from . import urls
 
-from playwright.sync_api import sync_playwright, TimeoutError
+import asyncio
+from playwright.async_api import async_playwright, TimeoutError
+from playwright_stealth import stealth_async
 from requests.cookies import cookiejar_from_dict
-from playwright_stealth import stealth_sync
 
 
 # Constants
@@ -14,67 +15,26 @@
 
 class SessionManager:
     def __init__(self) -> None:
+        """ 
+        This class is using asynchonous playwright mode. 
+        """
         self.headers = None
         self.session = requests.Session()
-
-        self.playwright = sync_playwright().start()
-        if self.browserType == "firefox":
-            self.browser = self.playwright.firefox.launch(
-                headless=self.headless
-            )
-        else:
-            #webkit doesn't or no longer works when trying to log in.
-            raise ValueError("Only supported browserType is 'firefox'")
-
-        user_agent = USER_AGENT + self.browser.version
-        self.page = self.browser.new_page(
-            user_agent=user_agent,
-            viewport=VIEWPORT
-        )
-
-        stealth_sync(self.page)
+        self.playwright = None
+        self.browser = None
+        self.page = None
 
     def check_auth(self):
         r = self.session.get(urls.account_info_v2())
         if r.status_code != 200:
             return False
         return True
 
-    def save_and_close_session(self):
-        cookies = {cookie["name"]: cookie["value"] for cookie in self.page.context.cookies()}
-        self.session.cookies = cookiejar_from_dict(cookies)
-        self.page.close()
-        self.browser.close()
-        self.playwright.stop()
-
     def get_session(self):
         return self.session
-
-    def sms_login(self, code):
-
-        # Inconsistent UI for SMS Authentication means we try both
-        try:
-            self.page.click("input[type=\"text\"]")
-            self.page.fill("input[type=\"text\"]", str(code))
-            self.page.click("text=Trust this device and skip this step in the future.")
-            with self.page.expect_navigation():
-                self.page.click("text=Log In")
-        except:
-            self.page.check("input[name=\"TrustDeviceChecked\"]")
-            self.page.click("[placeholder=\"Access Code\"]")
-            self.page.fill("[placeholder=\"Access Code\"]", str(code))
-            with self.page.expect_navigation():
-                self.page.click("text=Continue")
-
-        self.save_and_close_session()
-        return self.page.url == urls.account_summary()
-
-    def captureAuthToken(self, route):
-            self.headers = route.request.all_headers()
-            route.continue_()
 
     def login(self, username, password, totp_secret=None):
-        """ This function will log the user into schwab using Playwright and saving
+        """ This function will log the user into schwab using asynchronous Playwright and saving
         the authentication cookies in the session header. 
         :type username: str
         :param username: The username for the schwab account.
@@ -88,62 +48,67 @@ def login(self, username, password, totp_secret=None):
 
         :rtype: boolean
         :returns: True if login was successful and no further action is needed or False
-            if login requires additional steps (i.e. SMS)
+            if login requires additional steps (i.e. SMS - no longer supported)
         """
-
-        # Log in to schwab using Playwright
-        with self.page.expect_navigation():
-            self.page.goto("https://www.schwab.com/")
+        result = asyncio.run(self._async_login(username, password, totp_secret))
+        return result
+
+    async def _async_login(self, username, password, totp_secret=None):
+        """ This function runs in async mode to perform login.
+        Use with login function. See login function for details.
+        """
+        self.playwright = await async_playwright().start()
+        if self.browserType == "firefox":
+            self.browser = await self.playwright.firefox.launch(
+                headless=self.headless
+            )
+        else:
+            raise ValueError("Only supported browserType is 'firefox'")
 
+        user_agent = USER_AGENT + self.browser.version
+        self.page = await self.browser.new_page(
+            user_agent=user_agent,
+            viewport=VIEWPORT
+        )
+        await stealth_async(self.page)
+
+        await self.page.goto("https://www.schwab.com/")
 
-        # Capture authorization token.
-        self.page.route(re.compile(r".*balancespositions*"), self.captureAuthToken)
+        await self.page.route(re.compile(r".*balancespositions*"), self._asyncCaptureAuthToken)
 
-        # Wait for the login frame to load
         login_frame = "schwablmslogin"
-        self.page.wait_for_selector("#" + login_frame)
+        await self.page.wait_for_selector("#" + login_frame)
 
-        self.page.frame(name=login_frame).select_option("select#landingPageOptions", index=3)
+        await self.page.frame(name=login_frame).select_option("select#landingPageOptions", index=3)
+
+        await self.page.frame(name=login_frame).click("[placeholder=\"Login ID\"]")
+        await self.page.frame(name=login_frame).fill("[placeholder=\"Login ID\"]", username)
 
-        # Fill username
-        self.page.frame(name=login_frame).click("[placeholder=\"Login ID\"]")
-        self.page.frame(name=login_frame).fill("[placeholder=\"Login ID\"]", username)
-
-        # Add TOTP to password
         if totp_secret is not None:
             totp = pyotp.TOTP(totp_secret)
             password += str(totp.now())
 
-        # Fill password
-        self.page.frame(name=login_frame).press("[placeholder=\"Login ID\"]", "Tab")
-        self.page.frame(name=login_frame).fill("[placeholder=\"Password\"]", password)
+        await self.page.frame(name=login_frame).press("[placeholder=\"Login ID\"]", "Tab")
+        await self.page.frame(name=login_frame).fill("[placeholder=\"Password\"]", password)
 
-        # Submit
         try:
-            with self.page.expect_navigation():
-                self.page.frame(name=login_frame).press("[placeholder=\"Password\"]", "Enter")
+            await self.page.frame(name=login_frame).press("[placeholder=\"Password\"]", "Enter")
+            await self.page.wait_for_url(re.compile(r"app/trade"), wait_until="domcontentloaded") # Making it more robust than specifying an exact url which may change.
         except TimeoutError:
             raise Exception("Login was not successful; please check username and password")
 
-        # NOTE: THIS FUNCTIONALITY WILL SOON BE UNSUPPORTED/DEPRECATED.
-        if self.page.url != urls.trade_ticket():
-            # We need further authentication, so we'll send an SMS
-            print("Authentication state is not available. We will need to go through two factor authentication.")
-            print("You should receive a code through SMS soon")
-
-            # Send an SMS. The UI is inconsistent so we'll try both.
-            try:
-                with self.page.expect_navigation():
-                    self.page.click("[aria-label=\"Text me a 6 digit security code\"]")
-            except:
-                self.page.click("input[name=\"DeliveryMethodSelection\"]")
-                self.page.click("text=Text Message")
-                self.page.click("input:has-text(\"Continue\")")
-            return False
-
-        self.page.wait_for_selector("#_txtSymbol")
-
-        # Save our session
-        self.save_and_close_session()
+        await self.page.wait_for_selector("#_txtSymbol")
 
+        await self._async_save_and_close_session()
         return True
+
+    async def _async_save_and_close_session(self):
+        cookies = {cookie["name"]: cookie["value"] for cookie in await self.page.context.cookies()}
+        self.session.cookies = cookiejar_from_dict(cookies)
+        await self.page.close()
+        await self.browser.close()
+        await self.playwright.stop()
+
+    async def _asyncCaptureAuthToken(self, route):
+        self.headers = await route.request.all_headers()
+        await route.continue_()