Skip to content

Commit 3d00961

Browse files
neo773neo773
committed
Add support for Shorts
1 parent 2997ed0 commit 3d00961

4 files changed

Lines changed: 118 additions & 24 deletions

File tree

‎README.md‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,17 +74,26 @@ print(f"Video: {video}")
7474
channel = supadata.youtube.channel(id="https://youtube.com/@RickAstleyVEVO") # can be url, channel id, handle
7575
print(f"Channel: {channel}")
7676

77-
# Get a list of the channel video IDs
78-
channel_videos = supadata.youtube.channel.videos(id="RickAstleyVEVO") # can be url, channel id, or handle
79-
print(f"Channel Video IDs: {channel_videos}")
77+
# Get video IDs from a YouTube channel
78+
channel_videos = supadata.youtube.channel.videos(
79+
id="RickAstleyVEVO", # can be url, channel id, or handle
80+
type="all", # 'all', 'video', or 'short'
81+
limit=50
82+
)
83+
print(f"Regular videos: {channel_videos.video_ids}")
84+
print(f"Shorts: {channel_videos.short_ids}")
8085

8186
# Get Playlist metadata
8287
playlist = supadata.youtube.playlist(id="PLlaN88a7y2_plecYoJxvRFTLHVbIVAOoc") # can be url or playlist id
8388
print(f"Playlist: {playlist}")
8489

85-
# Get a list of the playlist video IDs
86-
playlist_videos = supadata.youtube.playlist.videos(id="https://www.youtube.com/playlist?list=PLlaN88a7y2_plecYoJxvRFTLHVbIVAOoc") # can be url or playlist id
87-
print(f"Playlist Videos IDs: {playlist_videos}")
90+
# Get video IDs from a YouTube playlist
91+
playlist_videos = supadata.youtube.playlist.videos(
92+
id="https://www.youtube.com/playlist?list=PLlaN88a7y2_plecYoJxvRFTLHVbIVAOoc", # can be url or playlist id
93+
limit=50
94+
)
95+
print(f"Regular videos: {playlist_videos.video_ids}")
96+
print(f"Shorts: {playlist_videos.short_ids}")
8897
```
8998

9099
## Error Handling

‎supadata/types.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,3 +254,22 @@ def __post_init__(self):
254254
self.last_updated = datetime.now()
255255
if self.channel is None:
256256
self.channel = YoutubeChannelBaseDict(id="", name="")
257+
258+
259+
@dataclass
260+
class VideoIds:
261+
"""Container for YouTube video IDs.
262+
263+
Attributes:
264+
video_ids: List of regular YouTube video IDs
265+
short_ids: List of YouTube Shorts IDs
266+
"""
267+
268+
video_ids: List[str] = None
269+
short_ids: List[str] = None
270+
271+
def __post_init__(self):
272+
if self.video_ids is None:
273+
self.video_ids = []
274+
if self.short_ids is None:
275+
self.short_ids = []

‎supadata/youtube.py‎

Lines changed: 32 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
"""YouTube-related operations for Supadata."""
22

33
from datetime import datetime
4-
from typing import Any, Callable, Dict, List
4+
from typing import Any, Callable, Dict, List, Optional, Literal
55

66
from .errors import SupadataError
77
from .types import (
@@ -11,6 +11,7 @@
1111
YoutubeChannel,
1212
YoutubePlaylist,
1313
YoutubeVideo,
14+
VideoIds,
1415
)
1516

1617

@@ -231,29 +232,42 @@ def __call__(self, id: str) -> YoutubeChannel:
231232

232233
return YoutubeChannel(**response)
233234

234-
def videos(self, id: str, limit: int | None = None) -> List[str]:
235-
"""Get a list of video IDs from a YouTube channel.
235+
def videos(
236+
self,
237+
id: str,
238+
limit: Optional[int] = None,
239+
type: Literal["all", "video", "short"] = "all"
240+
) -> VideoIds:
241+
"""Get video IDs from a YouTube channel.
236242
237243
Args:
238244
id: YouTube Channel ID
239245
limit: The limit of videos to be returned. None will
240246
return the default (30 videos)
247+
type: The type of videos to fetch.
248+
'all': Both regular videos and shorts (default)
249+
'video': Only regular videos
250+
'short': Only shorts
241251
242252
Returns:
243-
A list of video IDs.
253+
VideoIds object containing lists of video IDs and short IDs
244254
245255
Raises:
246256
SupadataError: If the API request fails
247257
"""
248258
self._youtube._validate_limit(limit)
249-
query_params = {"id": id}
259+
query_params = {"id": id, "type": type}
250260
if limit:
251261
query_params["limit"] = limit
252262

253263
response: dict = self._youtube._request(
254264
"GET", "/youtube/channel/videos", params=query_params
255265
)
256-
return response.get("video_ids", [])
266+
267+
return VideoIds(
268+
video_ids=response.get("video_ids", []),
269+
short_ids=response.get("short_ids", [])
270+
)
257271

258272
class _Playlist:
259273
def __init__(self, youtube: "YouTube"):
@@ -300,16 +314,20 @@ def __call__(self, id: str) -> YoutubePlaylist:
300314

301315
return YoutubePlaylist(**response, last_updated=last_updated)
302316

303-
def videos(self, id: str, limit: int | None = None) -> List[str]:
304-
"""Get a list of the IDs of the list of video IDs from a YouTube playlist.
317+
def videos(
318+
self,
319+
id: str,
320+
limit: Optional[int] = None
321+
) -> VideoIds:
322+
"""Get video IDs from a YouTube playlist.
305323
306324
Args:
307325
id: YouTube Playlist ID
308326
limit: The limit of videos to be returned. None will
309327
return the default (30 videos)
310328
311329
Returns:
312-
A list of video IDs.
330+
VideoIds object containing lists of video IDs and short IDs
313331
314332
Raises:
315333
SupadataError: If the API request fails
@@ -322,4 +340,8 @@ def videos(self, id: str, limit: int | None = None) -> List[str]:
322340
response: dict = self._youtube._request(
323341
"GET", "/youtube/playlist/videos", params=query_params
324342
)
325-
return response.get("video_ids", [])
343+
344+
return VideoIds(
345+
video_ids=response.get("video_ids", []),
346+
short_ids=response.get("short_ids", [])
347+
)

‎tests/test_client.py‎

Lines changed: 52 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -427,16 +427,50 @@ def test_youtube_channel_videos(client: Supadata, requests_mock) -> None:
427427
"videoIds": [
428428
"PQ2WjtaPfXU",
429429
"UIVADiGfwWc",
430+
],
431+
"shortIds": [
432+
"abc123",
433+
"def456",
430434
]
431435
}
432436

433437
requests_mock.get(
434-
f"{client.base_url}/youtube/channel/videos?id={channel_id}", json=mock_response
438+
f"{client.base_url}/youtube/channel/videos?id={channel_id}&type=all", json=mock_response
435439
)
436440
channel_videos = client.youtube.channel.videos(channel_id)
437-
assert isinstance(channel_videos, list)
438-
assert len(channel_videos) == len(mock_response["videoIds"])
439-
for i in channel_videos:
441+
assert hasattr(channel_videos, "video_ids")
442+
assert hasattr(channel_videos, "short_ids")
443+
assert isinstance(channel_videos.video_ids, list)
444+
assert isinstance(channel_videos.short_ids, list)
445+
assert len(channel_videos.video_ids) == len(mock_response["videoIds"])
446+
assert len(channel_videos.short_ids) == len(mock_response["shortIds"])
447+
for i in channel_videos.video_ids:
448+
assert i in mock_response["videoIds"]
449+
for i in channel_videos.short_ids:
450+
assert i in mock_response["shortIds"]
451+
452+
453+
def test_youtube_channel_videos_with_type(client: Supadata, requests_mock) -> None:
454+
channel_id = "UCsBjURrPoezyLs9EqgamOA"
455+
mock_response = {
456+
"videoIds": [
457+
"PQ2WjtaPfXU",
458+
"UIVADiGfwWc",
459+
],
460+
"shortIds": []
461+
}
462+
463+
requests_mock.get(
464+
f"{client.base_url}/youtube/channel/videos?id={channel_id}&type=video", json=mock_response
465+
)
466+
channel_videos = client.youtube.channel.videos(channel_id, type="video")
467+
assert hasattr(channel_videos, "video_ids")
468+
assert hasattr(channel_videos, "short_ids")
469+
assert isinstance(channel_videos.video_ids, list)
470+
assert isinstance(channel_videos.short_ids, list)
471+
assert len(channel_videos.video_ids) == len(mock_response["videoIds"])
472+
assert len(channel_videos.short_ids) == 0
473+
for i in channel_videos.video_ids:
440474
assert i in mock_response["videoIds"]
441475

442476

@@ -448,7 +482,7 @@ def test_youtube_channel_videos_invalid_id(client: Supadata, requests_mock) -> N
448482
"details": "The requested item could not be found.",
449483
}
450484
requests_mock.get(
451-
f"{client.base_url}/youtube/channel/videos?id={channel_id}",
485+
f"{client.base_url}/youtube/channel/videos?id={channel_id}&type=all",
452486
status_code=404,
453487
json=mock_response,
454488
)
@@ -468,6 +502,10 @@ def test_youtube_playlist_videos(client: Supadata, requests_mock) -> None:
468502
"videoIds": [
469503
"zDNaUi2cjv4",
470504
"B1t4Fjlomi8",
505+
],
506+
"shortIds": [
507+
"short1",
508+
"short2"
471509
]
472510
}
473511
requests_mock.get(
@@ -476,10 +514,16 @@ def test_youtube_playlist_videos(client: Supadata, requests_mock) -> None:
476514
)
477515

478516
playlist_videos = client.youtube.playlist.videos(playlist_id)
479-
assert isinstance(playlist_videos, list)
480-
assert len(playlist_videos) == len(mock_response["videoIds"])
481-
for i in playlist_videos:
517+
assert hasattr(playlist_videos, "video_ids")
518+
assert hasattr(playlist_videos, "short_ids")
519+
assert isinstance(playlist_videos.video_ids, list)
520+
assert isinstance(playlist_videos.short_ids, list)
521+
assert len(playlist_videos.video_ids) == len(mock_response["videoIds"])
522+
assert len(playlist_videos.short_ids) == len(mock_response["shortIds"])
523+
for i in playlist_videos.video_ids:
482524
assert i in mock_response["videoIds"]
525+
for i in playlist_videos.short_ids:
526+
assert i in mock_response["shortIds"]
483527

484528

485529
def test_youtube_playlist_videos_invalid_id(client: Supadata, requests_mock) -> None:

0 commit comments

Comments
 (0)