Update to get crumb when setting up session

dpguthrie · Jul 18, 2023 · 1ca552c · 1ca552c
1 parent c3d2641
commit 1ca552c
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 22 deletions.
diff --git a/yahooquery/base.py b/yahooquery/base.py
@@ -928,13 +928,14 @@ class _YahooFinance(object):
     def __init__(self, **kwargs):
         self.country = kwargs.get("country", "united states").lower()
         self.formatted = kwargs.pop("formatted", False)
-        self.session = _init_session(kwargs.pop("session", None), **kwargs)
+        self.session, self.crumb = _init_session(kwargs.pop("session", None), **kwargs)
         self.progress = kwargs.pop("progress", False)
         username = os.getenv("YF_USERNAME") or kwargs.get("username")
         password = os.getenv("YF_PASSWORD") or kwargs.get("password")
-        if username is not None and password is not None:
+        if username and password:
             self.login(username, password)
-        self.crumb = self.get_crumb()
+        if self.crumb is None:
+            self.set_crumb()
 
     @property
     def symbols(self):
@@ -947,24 +948,13 @@ def symbols(self):
     def symbols(self, symbols):
         self._symbols = _convert_to_list(symbols)
 
-    def get_crumb(self):
+    def set_crumb(self):
         response = self.session.get('https://query2.finance.yahoo.com/v1/test/getcrumb')
         if isinstance(self.session, FuturesSession):
             response = response.result()
         crumb = response.text
         if crumb is not None and len(crumb) > 0:
-            return crumb
-
-        response = self.session.get('https://finance.yahoo.com')
-        if isinstance(self.session, FuturesSession):
-            response = response.result()
-        path = re.compile(r'window\.YAHOO\.context = ({.*?});', re.DOTALL)
-        match = re.search(path, response.text)
-        if match:
-            js_dict = json.loads(match.group(1))
-            return js_dict.get('crumb', None)
-
-        return None
+            self.crumb = crumb
 
     @property
     def country(self):

diff --git a/yahooquery/misc.py b/yahooquery/misc.py
@@ -24,7 +24,9 @@ def _make_request(
                     country, ", ".join(sorted(COUNTRIES.keys()))
                 )
             )
-    session = _init_session(**kwargs)
+    session, crumb = _init_session(**kwargs)
+    if crumb is not None:
+        params['crumb'] = crumb
     r = getattr(session, method)(url, params=params, json=data)
     json = r.json()
     if response_field:

diff --git a/yahooquery/utils/__init__.py b/yahooquery/utils/__init__.py
@@ -1,4 +1,5 @@
 import datetime
+import json
 import random
 import re
 
@@ -139,6 +140,7 @@ def send(self, request, **kwargs):
 
 
 def _init_session(session=None, **kwargs):
+    crumb = None
     if session is None:
         if kwargs.get("asynchronous"):
             session = FuturesSession(max_workers=kwargs.get("max_workers", 8))
@@ -159,19 +161,29 @@ def _init_session(session=None, **kwargs):
                 max_retries=retries, timeout=kwargs.get("timeout", DEFAULT_TIMEOUT)
             ),
         )
-        session = setup_session_with_cookies(session)
-    return session
+        session, crumb = setup_session_with_cookies_and_crumb(session)
+    return session, crumb
 
 
-def setup_session_with_cookies(session: Session):
+def setup_session_with_cookies_and_crumb(session: Session):
     headers = {**random.choice(HEADERS), **addl_headers}
     session.headers = headers
-    response = session.get('https://finance.yahoo.com')
+    response = session.get('https://finance.yahoo.com', hooks={'response': get_crumb})
     if isinstance(session, FuturesSession):
         response = response.result()
-    return session
+    return session, response.crumb
 
 
+def get_crumb(r: Response, *args, **kwargs):
+    r.crumb = None
+    path = re.compile(r'window\.YAHOO\.context = ({.*?});', re.DOTALL)
+    match = re.search(path, r.text)
+    if match:
+        js_dict = json.loads(match.group(1))
+        r.crumb = js_dict.get('crumb', None)
+
+    return r
+
 def _flatten_list(ls):
     return [item for sublist in ls for item in sublist]