diff --git a/pyproject.toml b/pyproject.toml index 14b2eae4..958fb5b3 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/youtubeaio/youtube.py b/src/youtubeaio/youtube.py index 61f01377..83ff6df4 100644 --- a/src/youtubeaio/youtube.py +++ b/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: diff --git a/tests/test_short.py b/tests/test_short.py new file mode 100644 index 00000000..c20e5c9e --- /dev/null +++ b/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()