Add is_short method (#580)

Co-authored-by: Joost Lekkerkerker <[email protected]>

Author: electricsteve

pyproject.toml

@@ -53,6 +53,7 @@ Changelog = "https://github.com/joostlek/python-youtube/releases"
 
 [tool.coverage.report]
 show_missing = true
+fail_under = 80
 
 [tool.coverage.run]
 plugins = ["covdefaults"]

src/youtubeaio/youtube.py

@@ -77,6 +77,7 @@ def __init__(
             ],
         ] = {
             "get": self._api_get_request,
+            "head": self._api_head_request,
         }
 
     async def _check_request_return(self, response: ClientResponse) -> ClientResponse:
@@ -115,6 +116,18 @@ async def _api_get_request(
         response = await session.get(url, headers=headers, json=data)
         return await self._check_request_return(response)
 
+    async def _api_head_request(
+        self,
+        session: ClientSession,
+        url: str,
+        data: dict[str, Any] | None = None,
+    ) -> ClientResponse:
+        """Make HEAD request with authorization."""
+        headers = {"Authorization": f"Bearer {self._user_auth_token}"}
+        self.logger.debug("making HEAD request to %s", url)
+        response = await session.head(url, headers=headers, json=data)
+        return await self._check_request_return(response)
+
     async def _build_generator(
         self,
         req: str,
@@ -284,6 +297,20 @@ async def get_playlist_items(
         ):
             yield item  # type: ignore[misc]
 
+    async def is_short(self, video_id: str) -> bool:
+        """Return True if the video ID corresponds to a YouTube Short."""
+        _url = "https://www.youtube.com/shorts/" + video_id
+        if not self.session:
+            self.session = ClientSession()
+            self._close_session = True
+        async with asyncio.timeout(self.session_timeout):
+            response = await self._api_head_request(self.session, _url)
+        if response.status == 200:
+            return True
+        if response.status == 303:
+            return False
+        raise YouTubeAPIError
+
     async def close(self) -> None:
         """Close open client session."""
         if self.session and self._close_session:

tests/test_short.py

@@ -0,0 +1,49 @@
+"""Tests for the YouTube client."""
+
+import aiohttp
+from aresponses import ResponsesMockServer
+
+from youtubeaio.youtube import YouTube
+
+
+async def test_is_short(
+    aresponses: ResponsesMockServer,
+) -> None:
+    """Test YouTube video id that is a short."""
+    aresponses.add(
+        "www.youtube.com",
+        "/shorts/2YUPfsi8PF4",
+        "HEAD",
+        aresponses.Response(
+            status=200,
+            headers={"Content-Type": "text/html"},
+        ),
+    )
+    async with aiohttp.ClientSession() as session:
+        youtube = YouTube(session=session)
+        response = await youtube.is_short("2YUPfsi8PF4")
+        assert response
+        await youtube.close()
+
+
+async def test_is_not_short(
+    aresponses: ResponsesMockServer,
+) -> None:
+    """Test YouTube video id that is not a short."""
+    aresponses.add(
+        "www.youtube.com",
+        "/shorts/KXaMtA6kWXU",
+        "HEAD",
+        aresponses.Response(
+            status=303,
+            headers={
+                "Content-Type": "application/binary",
+                "Location": "https://www.youtube.com/watch?v=KXaMtA6kWXU",
+            },
+        ),
+    )
+    async with aiohttp.ClientSession() as session:
+        youtube = YouTube(session=session)
+        response = await youtube.is_short("KXaMtA6kWXU")
+        assert not response
+        await youtube.close()