B

API/app.py

@@ -225,33 +225,28 @@ async def download_image(session, post_id, post_url):
         content = await response.read()
         return (post_id, content)
 
-async def tar_stream_generator(posts, session):
-    for post_id, image_url in posts.items():
-        try:
-            # Download the image data
-            image_data = await download_image(session, post_id, image_url)
-            # Create the tar header
-            image_name = f"{post_id}.png"
-            tarinfo = tarfile.TarInfo(name=image_name)
-            tarinfo.size = len(image_data)
-            tarinfo.mtime = int(time.time())
-            tarinfo.mode = 0o644
-            tarinfo.uid = 0
-            tarinfo.gid = 0
-            header = tarinfo.tobuf()
-            # Yield the header
-            yield header
-            # Yield the image data
-            yield image_data
-            # Pad the image data to a multiple of 512 bytes
-            remainder = len(image_data) % 512
-            if remainder > 0:
-                padding = b'\0' * (512 - remainder)
-                yield padding
-        except Exception as e:
-            continue
-    # At the end, yield two 512-byte blocks of zeros
-    yield b'\0' * 1024
+async def generate_tar(posts, session):
+    tar_buffer = io.BytesIO()
+    tarf = tarfile.open(mode="w", fileobj=tar_buffer)
+
+    semaphore = asyncio.Semaphore(100)
+
+    async def add_to_tar(post_id, image_url):
+        async with semaphore:
+            url = f"{img_base}{image_url}"
+            async with session.get(url) as response:
+                image_data = await response.read()
+                image_name = f"{post_id}.png"
+                info = tarfile.TarInfo(name=image_name)
+                info.size = len(image_data)
+                tarf.addfile(tarinfo=info, fileobj=io.BytesIO(image_data))
+
+    tasks = [add_to_tar(post_id, image_url) for post_id, image_url in posts.items()]
+    await asyncio.gather(*tasks)
+
+    tarf.close()
+    tar_buffer.seek(0)
+    return tar_buffer
 
 @app.post("/download")
 async def download(
@@ -259,13 +254,13 @@ async def download(
 ):
     posts = items.posts
     async with aiohttp.ClientSession(cookies=cookies, headers=headers) as session:
-        tar_generator = tar_stream_generator(posts, session)
+        tar_buffer = await generate_tar(posts, session)
 
-        return StreamingResponse(
-            tar_generator,
-            media_type="application/x-tar",
-            headers={"Content-Disposition": f"attachment; filename={base26_time()}.tar"}
-        )
+    return StreamingResponse(
+        tar_buffer,
+        media_type="application/x-tar",
+        headers={"Content-Disposition": f"attachment; filename={base26_time()}.tar"}
+    )
 
 @app.get("/")
 async def read_root():