Add the ability to use YouTube API

This is to circumvent YouTube blocking.

2 files changed, 18 insertions, 10 deletions

diff --git a/data_analyser/__main__.py b/data_analyser/__main__.py
index 8b4a984..ad85c09 100755
--- a/data_analyser/__main__.py
+++ b/data_analyser/__main__.py
@@ -1,10 +1,15 @@
 #!/usr/bin/env python
 
+from os import getenv
 from time import sleep
-from .utils import analyse_collection
+
+from dotenv import load_dotenv
 
 from components.database import subscriptions
+from .utils import analyse_collection
+
+load_dotenv('.env')
 
 while True:
-    analyse_collection(subscriptions)
+    analyse_collection(subscriptions, getenv("YOUTUBE_API_KEY") or '')
     sleep(30)
diff --git a/data_analyser/utils.py b/data_analyser/utils.py
index 95dc2fe..921015d 100644
--- a/data_analyser/utils.py
+++ b/data_analyser/utils.py
@@ -1,31 +1,34 @@
+from traceback import print_exc
+
+from pymongo.collection import Collection
+
 from components.subscriptions.main import Subscription
 from components.subscriptions.typing import SubsDict
 from components.videos import VideoTuple
 from components.extractor.obtain_vid_info import obtain_vid_duration
 
-from pymongo.collection import Collection
-
-def analyse_video(vid_tuple: VideoTuple) -> VideoTuple:
+def analyse_video(vid_tuple: VideoTuple, api_key: str='') -> VideoTuple:
     try:
-        duration = obtain_vid_duration(vid_tuple.link)
+        duration = obtain_vid_duration(vid_tuple.link, vid_tuple.id, api_key=api_key)
     except:
+        print_exc()
         duration = -2
     return vid_tuple._replace(analysed=True, duration=duration)
 
-def analyse_subscription(sub: Subscription) -> bool:
+def analyse_subscription(sub: Subscription, api_key: str='') -> bool:
     updated = False
     for i, vid in enumerate(sub.videos):
         if not vid.analysed:
-            sub.videos[i] = analyse_video(vid)
+            sub.videos[i] = analyse_video(vid, api_key)
             updated = True
     return updated
 
-def analyse_collection(subs_collection: Collection[SubsDict]) -> int:
+def analyse_collection(subs_collection: Collection[SubsDict], api_key: str='') -> int:
     num_updated = 0
     for sub_dict in subs_collection.find():
         sub = Subscription(**sub_dict)
         sub._collection = subs_collection
-        if analyse_subscription(sub):
+        if analyse_subscription(sub, api_key):
             sub.update_videos()
             num_updated += 1
     return num_updated