From 638122081b178a6c7990c2a811238b344d292964 Mon Sep 17 00:00:00 2001
From: cmooredev <cmooredev@pm.me>
Date: Fri, 16 Aug 2024 06:40:37 -0700
Subject: [PATCH] Add streaming download option to Video.bytes()

---
 TikTokApi/api/video.py | 23 ++++++++++++++++++-----
 requirements.txt       |  1 +
 setup.py               |  2 +-
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/TikTokApi/api/video.py b/TikTokApi/api/video.py
index 16271c4b..992457d2 100644
--- a/TikTokApi/api/video.py
+++ b/TikTokApi/api/video.py
@@ -5,6 +5,8 @@
 import requests
 from ..exceptions import InvalidResponseException
 import json
+import httpx
+from typing import Union, AsyncIterator
 
 if TYPE_CHECKING:
     from ..tiktok import TikTokApi
@@ -162,7 +164,7 @@ async def info(self, **kwargs) -> dict:
         )
         return video_info
 
-    async def bytes(self, **kwargs) -> bytes:
+    async def bytes(self, stream: bool = False, **kwargs) -> Union[bytes, AsyncIterator[bytes]]:
         """
         Returns the bytes of a TikTok Video.
 
@@ -172,13 +174,16 @@ async def bytes(self, **kwargs) -> bytes:
         Example Usage:
             .. code-block:: python
 
-                video_bytes = api.video(id='7041997751718137094').bytes()
+                video_bytes = await api.video(id='7041997751718137094').bytes()
 
                 # Saving The Video
                 with open('saved_video.mp4', 'wb') as output:
                     output.write(video_bytes)
-        """
 
+                # Streaming (if stream=True)
+                async for chunk in api.video(id='7041997751718137094').bytes(stream=True):
+                    # Process or upload chunk
+        """
         i, session = self.parent._get_session(**kwargs)
         downloadAddr = self.as_dict["video"]["downloadAddr"]
 
@@ -189,8 +194,16 @@ async def bytes(self, **kwargs) -> bytes:
         h["accept-encoding"] = 'identity;q=1, *;q=0'
         h["referer"] = 'https://www.tiktok.com/'
 
-        resp = requests.get(downloadAddr, headers=h, cookies=cookies)
-        return resp.content
+        if stream:
+            async def stream_bytes():
+                async with httpx.AsyncClient() as client:
+                    async with client.stream('GET', downloadAddr, headers=h, cookies=cookies) as response:
+                        async for chunk in response.aiter_bytes():
+                            yield chunk
+            return stream_bytes()
+        else:
+            resp = requests.get(downloadAddr, headers=h, cookies=cookies)
+            return resp.content
 
     def __extract_from_data(self) -> None:
         data = self.as_dict
diff --git a/requirements.txt b/requirements.txt
index fe477b5f..ea53cd3a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 requests>=2.31.0,<3.0
 playwright>=1.36.0,<2.0
+httpx>=0.27.0,<1.0
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 6369da86..8d531683 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@
     long_description_content_type="text/markdown",
     download_url="https://github.com/davidteather/TikTok-Api/tarball/main",
     keywords=["tiktok", "python3", "api", "unofficial", "tiktok-api", "tiktok api"],
-    install_requires=["requests", "playwright"],
+    install_requires=["requests", "playwright", "httpx"],
     classifiers=[
         "Development Status :: 4 - Beta",
         "Intended Audience :: Developers",