From ca5cc9d80cc1843f6a0a5a0d20f3e12674d26b38 Mon Sep 17 00:00:00 2001
From: Russell Newton <russell.newton01@gmail.com>
Date: Sun, 19 Nov 2023 18:54:05 -0500
Subject: [PATCH] Added a video download function for nonslideshows

---
 docs/source/users/usage.rst     |  6 +++++
 pyproject.toml                  |  3 +++
 src/tiktokapipy/models/video.py | 40 +++++++++++++++++++++++++++++++--
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/docs/source/users/usage.rst b/docs/source/users/usage.rst
index 80fd660..30e6189 100644
--- a/docs/source/users/usage.rst
+++ b/docs/source/users/usage.rst
@@ -200,6 +200,12 @@ Download Videos and Slideshows
 If all you want to do is download a video or slideshow from TikTok, go no further. Because slideshows are saved as
 images with a sound, you'll need to join these images together with the sound. I'd suggest using `ffmpeg`_ for this:
 
+.. note::
+    The ``download()`` function on videos can be used to download videos to a file. It doesn't work for slideshows.
+    It requires the ``yt-dlp`` package to be installed, which can be installed with ``pip install yt-dlp`` or
+    ``pip install tiktokapipy[download]``.
+
+
 .. code-block:: py
 
     import asyncio
diff --git a/pyproject.toml b/pyproject.toml
index 20fc63b..f4ceda8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,6 +52,9 @@ docs = [
     "sphinx-tabs",
     "mock",
 ]
+download = [
+    "yt-dlp"
+]
 
 [project.urls]
 "Homepage" = "https://github.com/Russell-Newton/TikTokPy"
diff --git a/src/tiktokapipy/models/video.py b/src/tiktokapipy/models/video.py
index fd5d3c4..f0c9583 100644
--- a/src/tiktokapipy/models/video.py
+++ b/src/tiktokapipy/models/video.py
@@ -8,7 +8,7 @@
 from typing import Any, ForwardRef, List, Optional, Union
 
 from playwright.async_api import BrowserContext as AsyncBrowserContext
-from pydantic import Field, computed_field
+from pydantic import AliasChoices, Field, computed_field
 from tiktokapipy import TikTokAPIError
 from tiktokapipy.models import CamelCaseModel, TitleCaseModel
 from tiktokapipy.util.deferred_collectors import (
@@ -101,7 +101,7 @@ class ImagePost(CamelCaseModel):
 
 
 class LightVideo(CamelCaseModel):
-    id: int = Field(alias="id")
+    id: int = Field(validation_alias=AliasChoices("cid", "uid", "id"))
     stats: VideoStats
     create_time: datetime
 
@@ -115,6 +115,8 @@ class Video(LightVideo):
     digged: bool
     item_comment_status: int
     author: Union[LightUser, str]
+    image_post: Optional[ImagePost] = None
+    """The images in the video if the video is a slideshow"""
 
     @computed_field(repr=False)
     @property
@@ -170,6 +172,40 @@ def creator(self) -> Union[DeferredUserGetterAsync, DeferredUserGetterSync]:
     def url(self) -> str:
         return video_link(self.id)
 
+    def download(self) -> str:
+        """
+        Downloads this video, returning the relative filepath where it was stored.
+        Requires yt-dlp installed (``pip install yt-dlp``
+        or ``pip install tiktokapipy[download]``)
+        """
+        if self.image_post:
+            raise TikTokAPIError(
+                "The download function isn't available for slideshows."
+            )
+
+        try:
+            import yt_dlp
+        except ImportError:
+            raise TikTokAPIError(
+                "You don't have youtube_dl installed! "
+                "Please install with `pip install yt-dlp` or "
+                "`pip install tiktokapipy[download]"
+            )
+
+        downloaded_file = ""
+
+        class GetFileNamePP(yt_dlp.postprocessor.PostProcessor):
+            def run(self, info):
+                nonlocal downloaded_file
+                downloaded_file = info["filename"]
+                return [], info
+
+        with yt_dlp.YoutubeDL() as ydl:
+            ydl.add_post_processor(GetFileNamePP())
+            ydl.download([video_link(self.id)])
+
+        return downloaded_file
+
 
 del Challenge, LightChallenge, Comment, LightUser, User, UserStats
 from tiktokapipy.models.challenge import Challenge, LightChallenge  # noqa E402