Add playlist class (#1207)

* feat: add playlist class

TikTokApi/api/playlist.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, ClassVar, Iterator, Optional
+from ..exceptions import InvalidResponseException
+
+if TYPE_CHECKING:
+    from ..tiktok import TikTokApi
+    from .video import Video
+    from .user import User
+
+
+class Playlist:
+    """
+    A TikTok video playlist.
+
+    Example Usage:
+        .. code-block:: python
+
+            playlist = api.playlist(id='7426714779919797038')
+    """
+
+    parent: ClassVar[TikTokApi]
+
+    id: Optional[str]
+    """The ID of the playlist."""
+    name: Optional[str]
+    """The name of the playlist."""
+    video_count: Optional[int]
+    """The video count of the playlist."""
+    creator: Optional[User]
+    """The creator of the playlist."""
+    cover_url: Optional[str]
+    """The cover URL of the playlist."""
+    as_dict: dict
+    """The raw data associated with this Playlist."""
+
+    def __init__(
+        self,
+        id: Optional[str] = None,
+        data: Optional[dict] = None,
+    ):
+        """
+        You must provide the playlist id or playlist data otherwise this
+        will not function correctly.
+        """
+
+        if id is None and data.get("id") is None:
+            raise TypeError("You must provide id parameter.")
+
+        self.id = id
+
+        if data is not None:
+            self.as_dict = data
+            self.__extract_from_data()
+
+    async def info(self, **kwargs) -> dict:
+        """
+        Returns a dictionary of information associated with this Playlist.
+
+        Returns:
+            dict: A dictionary of information associated with this Playlist.
+
+        Raises:
+            InvalidResponseException: If TikTok returns an invalid response, or one we don't understand.
+
+        Example Usage:
+            .. code-block:: python
+
+                user_data = await api.playlist(id='7426714779919797038').info()
+        """
+
+        id = getattr(self, "id", None)
+        if not id:
+            raise TypeError(
+                "You must provide the playlist id when creating this class to use this method."
+            )
+
+        url_params = {
+            "mixId": id,
+            "msToken": kwargs.get("ms_token"),
+        }
+
+        resp = await self.parent.make_request(
+            url="https://www.tiktok.com/api/mix/detail/",
+            params=url_params,
+            headers=kwargs.get("headers"),
+            session_index=kwargs.get("session_index"),
+        )
+
+        if resp is None:
+            raise InvalidResponseException(resp, "TikTok returned an invalid response.")
+
+        self.as_dict = resp["mixInfo"]
+        self.__extract_from_data()
+        return resp
+
+    async def videos(self, count=30, cursor=0, **kwargs) -> Iterator[Video]:
+        """
+        Returns an iterator of videos in this User's playlist.
+
+        Returns:
+            Iterator[dict]: An iterator of videos in this User's playlist.
+
+        Raises:
+            InvalidResponseException: If TikTok returns an invalid response, or one we don't understand.
+
+        Example Usage:
+            .. code-block:: python
+
+                playlist_videos = await api.playlist(id='7426714779919797038').videos()
+        """
+        id = getattr(self, "id", None)
+        if id is None or id == "":
+            await self.info(**kwargs)
+
+        found = 0
+        while found < count:
+            params = {
+              "mixId": id,
+              "count": min(count, 30),
+              "cursor": cursor,
+          }
+
+            resp = await self.parent.make_request(
+                url="https://www.tiktok.com/api/mix/item_list/",
+                params=params,
+                headers=kwargs.get("headers"),
+                session_index=kwargs.get("session_index"),
+            )
+
+            if resp is None:
+                raise InvalidResponseException(
+                    resp, "TikTok returned an invalid response."
+                )
+
+            for video in resp.get("itemList", []):
+                yield self.parent.video(data=video)
+                found += 1
+
+            if not resp.get("hasMore", False):
+                return
+
+            cursor = resp.get("cursor")
+
+    def __extract_from_data(self):
+        data = self.as_dict
+        keys = data.keys()
+
+        if "mixInfo" in keys:
+            data = data["mixInfo"]
+
+        self.id = data.get("id", None) or data.get("mixId", None)
+        self.name = data.get("name", None) or data.get("mixName", None)
+        self.video_count = data.get("videoCount", None)
+        self.creator = self.parent.user(data=data.get("creator", {}))
+        self.cover_url = data.get("cover", None)
+
+        if None in [self.id, self.name, self.video_count, self.creator, self.cover_url]:
+            User.parent.logger.error(
+                f"Failed to create Playlist with data: {data}\nwhich has keys {data.keys()}"
+            )
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __str__(self):
+        id = getattr(self, "id", None)
+        return f"TikTokApi.playlist(id='{id}'')"

TikTokApi/api/user.py

@@ -5,6 +5,7 @@
 if TYPE_CHECKING:
     from ..tiktok import TikTokApi
     from .video import Video
+    from .playlist import Playlist
 
 
 class User:
@@ -87,20 +88,21 @@ async def info(self, **kwargs) -> dict:
         self.__extract_from_data()
         return resp
 
-    async def playlists(self, count=20, cursor=0, **kwargs) -> Iterator[dict]:
+    async def playlists(self, count=20, cursor=0, **kwargs) -> Iterator[Playlist]:
         """
-        Returns a dictionary of information associated with this User's playlist.
+        Returns a user's playlists.
 
         Returns:
-            dict: A dictionary of information associated with this User's playlist.
+            async iterator/generator: Yields TikTokApi.playlist objects.
 
         Raises:
             InvalidResponseException: If TikTok returns an invalid response, or one we don't understand.
 
         Example Usage:
             .. code-block:: python
 
-                user_data = await api.user(username='therock').playlist()
+                async for playlist in await api.user(username='therock').playlists():
+                    # do something
         """
 
         sec_uid = getattr(self, "sec_uid", None)
@@ -109,30 +111,30 @@ async def playlists(self, count=20, cursor=0, **kwargs) -> Iterator[dict]:
         found = 0
 
         while found < count:
-          params = {
-              "secUid": sec_uid,
-              "count": 20,
-              "cursor": cursor,
-          }
-
-          resp = await self.parent.make_request(
-              url="https://www.tiktok.com/api/user/playlist",
-              params=params,
-              headers=kwargs.get("headers"),
-              session_index=kwargs.get("session_index"),
-          )
-
-          if resp is None:
-              raise InvalidResponseException(resp, "TikTok returned an invalid response.")
-
-          for playlist in resp.get("playList", []):
-              yield playlist
-              found += 1
-
-          if not resp.get("hasMore", False):
-              return
-
-          cursor = resp.get("cursor")
+            params = {
+                "secUid": self.sec_uid,
+                "count": min(count, 20),
+                "cursor": cursor,
+            }
+
+            resp = await self.parent.make_request(
+                url="https://www.tiktok.com/api/user/playlist",
+                params=params,
+                headers=kwargs.get("headers"),
+                session_index=kwargs.get("session_index"),
+            )
+
+            if resp is None:
+                raise InvalidResponseException(resp, "TikTok returned an invalid response.")
+            
+            for playlist in resp.get("playList", []):
+                yield self.parent.playlist(data=playlist)
+                found += 1
+            
+            if not resp.get("hasMore", False):
+                return
+            
+            cursor = resp.get("cursor")
 
 
     async def videos(self, count=30, cursor=0, **kwargs) -> Iterator[Video]:

TikTokApi/tiktok.py

@@ -18,6 +18,7 @@
 from .api.comment import Comment
 from .api.trending import Trending
 from .api.search import Search
+from .api.playlist import Playlist
 
 from .exceptions import (
     InvalidJSONException,
@@ -55,6 +56,7 @@ class TikTokApi:
     comment = Comment
     trending = Trending
     search = Search
+    playlist = Playlist
 
     def __init__(self, logging_level: int = logging.WARN, logger_name: str = None):
         """
@@ -77,6 +79,7 @@ def __init__(self, logging_level: int = logging.WARN, logger_name: str = None):
         Comment.parent = self
         Trending.parent = self
         Search.parent = self
+        Playlist.parent = self
 
     def __create_logger(self, name: str, level: int = logging.DEBUG):
         """Create a logger for the class."""

examples/playlist_example.py

@@ -0,0 +1,24 @@
+from TikTokApi import TikTokApi
+import asyncio
+import os
+
+ms_token = os.environ.get(
+    "ms_token", None
+)  # set your own ms_token, think it might need to have visited a profile
+
+
+async def user_example():
+    async with TikTokApi() as api:
+        await api.create_sessions(ms_tokens=[ms_token], num_sessions=1, sleep_after=3)
+        user = api.user("therock")
+
+        async for playlist in user.playlists(count=3):
+            print(playlist)
+            print(playlist.name)
+
+            async for video in playlist.videos(count=3):
+                print(video)
+                print(video.url)
+
+if __name__ == "__main__":
+    asyncio.run(user_example())

tests/test_playlist.py

@@ -0,0 +1,38 @@
+from TikTokApi import TikTokApi
+import os
+import pytest
+
+playlist_id="7281443725770476321"
+playlist_name="Doctor Who"
+playlist_creator="bbc"
+
+ms_token = os.environ.get("ms_token", None)
+
+
+@pytest.mark.asyncio
+async def test_playlist_info():
+    api = TikTokApi()
+    async with api:
+        await api.create_sessions(ms_tokens=[ms_token], num_sessions=1, sleep_after=3)
+        playlist = api.playlist(id=playlist_id)
+        await playlist.info()
+
+        assert playlist.id == playlist_id
+        assert playlist.name == playlist_name
+        assert playlist.creator.username == playlist_creator
+        assert playlist.video_count > 0
+        assert playlist.cover_url is not None
+        assert playlist.as_dict is not None
+
+@pytest.mark.asyncio
+async def test_playlist_videos():
+    api = TikTokApi()
+    async with api:
+        await api.create_sessions(ms_tokens=[ms_token], num_sessions=1, sleep_after=3)
+        playlist = api.playlist(id=playlist_id)
+
+        count = 0
+        async for video in playlist.videos(count=30):
+            count += 1
+
+        assert count >= 30