From d5885273ac39050fd360b78faa3d37e85fd3cf4e Mon Sep 17 00:00:00 2001
From: simon <simobilleter@gmail.com>
Date: Sat, 7 Jan 2023 15:41:11 +0700
Subject: [PATCH] detect any shorts before adding to queue

---
 tubearchivist/home/src/download/queue.py | 23 +++++++++++++++++++----
 tubearchivist/home/src/ta/helper.py      |  2 +-
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/tubearchivist/home/src/download/queue.py b/tubearchivist/home/src/download/queue.py
index 34afbbbf..9e62841c 100644
--- a/tubearchivist/home/src/download/queue.py
+++ b/tubearchivist/home/src/download/queue.py
@@ -17,7 +17,7 @@ from home.src.es.connect import ElasticWrap, IndexPaginate
 from home.src.index.playlist import YoutubePlaylist
 from home.src.index.video_constants import VideoTypeEnum
 from home.src.ta.config import AppConfig
-from home.src.ta.helper import DurationConverter
+from home.src.ta.helper import DurationConverter, is_shorts
 from home.src.ta.ta_redis import RedisArchivist
 
 
@@ -204,9 +204,8 @@ class PendingList(PendingIndex):
         video_results = playlist.json_data.get("playlist_entries")
         youtube_ids = [i["youtube_id"] for i in video_results]
         for video_id in youtube_ids:
-            # FIXME: This will need to be adjusted to support Live/Shorts
-            #  from playlists
-            self._add_video(video_id, VideoTypeEnum.VIDEOS)
+            # match vid_type later
+            self._add_video(video_id, VideoTypeEnum.UNKNOWN)
 
     def add_to_pending(self, status="pending"):
         """add missing videos to pending list"""
@@ -269,9 +268,25 @@ class PendingList(PendingIndex):
 
         if vid["live_status"] == "was_live":
             vid_type = VideoTypeEnum.STREAMS
+        else:
+            if self._check_shorts(vid):
+                vid_type = VideoTypeEnum.SHORTS
 
         return self._parse_youtube_details(vid, vid_type)
 
+    @staticmethod
+    def _check_shorts(vid):
+        """check if vid is shorts video"""
+        if vid["width"] > vid["height"]:
+            return False
+
+        duration = vid.get("duration")
+        if duration and isinstance(duration, int):
+            if duration > 60:
+                return False
+
+        return is_shorts(vid["id"])
+
     def _parse_youtube_details(self, vid, vid_type=VideoTypeEnum.VIDEOS):
         """parse response"""
         vid_id = vid.get("id")
diff --git a/tubearchivist/home/src/ta/helper.py b/tubearchivist/home/src/ta/helper.py
index 51bbf428..f07dc805 100644
--- a/tubearchivist/home/src/ta/helper.py
+++ b/tubearchivist/home/src/ta/helper.py
@@ -135,7 +135,7 @@ def get_mapping():
     return index_config
 
 
-def is_short(youtube_id):
+def is_shorts(youtube_id):
     """check if youtube_id is a shorts video, bot not it it's not a shorts"""
     shorts_url = f"https://www.youtube.com/shorts/{youtube_id}"
     response = requests.head(