Youtube shorts support

Pipfile

@@ -15,6 +15,7 @@ ffmpeg-python = "==0.2.0"
 opencv-python = "==4.8.0.74"
 asyncpraw = "==7.7.1"
 twscrape = "==0.7.0"
+pytube = "==15.0.0"
 
 [dev-packages]
 black = "==23.7.0"

README.md

@@ -10,6 +10,7 @@ A Discord bot that automatically embeds messages containing a video link of a su
 - Tiktok ✅
 - Reddit ✅
 - Twitter ✅
+- Youtube shorts ✅
 
 ## How to run
 - Build the docker image: `docker build . -t video-embed-bot` or simply pull it from ghcr:

downloader/registry.py

@@ -4,6 +4,7 @@
 from downloader import reddit
 from downloader import tiktok
 from downloader import twitter
+from downloader import youtube
 
 
 CLASSES = {
@@ -12,6 +13,7 @@
     facebook.FacebookClient,
     reddit.RedditClient,
     twitter.TwitterClient,
+    youtube.YoutubeClient,
 }
 
 

downloader/youtube.py

@@ -0,0 +1,28 @@
+import io
+
+import pytube
+
+from downloader import base
+from models import post
+
+
+class YoutubeClient(base.BaseClient):
+    DOMAINS = ['youtube.com/shorts']
+
+    async def get_post(self) -> post.Post:
+        vid = pytube.YouTube(self.url)
+
+        p = post.Post(
+            url=self.url,
+            author=vid.author,
+            description=vid.title,
+            views=vid.views,
+            created=vid.publish_date,
+            buffer=io.BytesIO(),
+        )
+
+        vid.streams.filter(progressive=True, file_extension='mp4').order_by(
+            'resolution'
+        ).desc().first().stream_to_buffer(p.buffer)
+
+        return p

models/post.py

@@ -26,16 +26,22 @@ def __str__(self) -> str:
         ).format(
             url=self.url,
             author=self.author or '❌',
-            created=self.created.strftime('%H:%M · %b %-d, %Y') if self.created else '❌',
+            created=self._date_human_format(date=self.created) if self.created else '❌',
             description=self.description or '❌',
-            views=self._human_format(self.views) if self.views else '❌',
-            likes=self._human_format(self.likes) if self.likes else '❌',
+            views=self._number_human_format(num=self.views) if self.views else '❌',
+            likes=self._number_human_format(num=self.likes) if self.likes else '❌',
         )
 
-    def _human_format(self, num: int) -> str:
+    def _number_human_format(self, num: int) -> str:
         num = float('{:.3g}'.format(num))
         magnitude = 0
         while abs(num) >= 1000:
             magnitude += 1
             num /= 1000.0
         return '{}{}'.format('{:f}'.format(num).rstrip('0').rstrip('.'), ['', 'K', 'M', 'B', 'T'][magnitude])
+
+    def _date_human_format(self, date: datetime.datetime) -> str:
+        if date.hour == 0 and date.minute == 0:
+            return date.strftime('%b %-d, %Y')
+
+        return date.strftime('%H:%M · %b %-d, %Y')

utils.py

@@ -14,4 +14,4 @@ def find_first_url(string: str) -> typing.Optional[str]:
 def guess_extension_from_buffer(buffer: io.BytesIO) -> str:
     extension = mimetypes.guess_extension(type=magic.from_buffer(buffer.read(2048), mime=True))
     buffer.seek(0)
-    return extension
+    return extension or '.mp4'