diff --git a/main3.py b/main3.py index e49203d..4873491 100644 --- a/main3.py +++ b/main3.py @@ -26,18 +26,12 @@ def save_to_checkpoint(path: Path): f.write(str(path) + "\n") -# def utc_filename(timestamp: str, url: str) -> Path: -# dt = datetime.fromisoformat(timestamp.replace("Z", "+00:00")) -# ext = "mp4" if ".mp4" in url.lower() else "jpg" -# return OUTPUT_DIR / f"{dt.strftime('%Y-%m-%d_%H-%M-%S')}.{ext}" - -def utc_filename(timestamp: str, url: str) -> Path: +async def utc_filename(timestamp: str, url: str) -> Path: # parse "YYYY-MM-DD HH:MM:SS UTC" dt = datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S UTC") dt = dt.replace(tzinfo=pytz.utc) ext = "mp4" if ".mp4" in url.lower() else "jpg" out_path = os.path.join("./downloads",f"{dt.strftime('%Y-%m-%d_%H-%M-%S')}.{ext}") - print(out_path) return Path(out_path) @@ -51,22 +45,24 @@ async def get_cdn_url(download_link: str) -> str: return response.text.strip() -async def download_one(session, sem, url, path, failures, stats): - if path.exists(): - return - - # Already downloaded? (checkpoint resume) - if str(path) in stats["done"]: - return - +async def download_one(session, sem, url, timestamp, failures, stats): async with sem: for attempt in range(1, RETRIES + 1): try: cdn_url = await get_cdn_url(url) - async with session.get(cdn_url) as resp: + url = cdn_url + path = await utc_filename(timestamp, url) + async with session.get(url) as resp: resp.raise_for_status() data = await resp.read() + if path.exists(): + return + + # Already downloaded? (checkpoint resume) + if str(path) in stats["done"]: + return + path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(data) @@ -89,15 +85,20 @@ async def main(): with open("./resources/json/memories_history.json", "r") as f: memories = json.load(f)["Saved Media"] - # Prepare tasks - tasks_to_download = [] - for item in memories: - url = item.get("Download Link") - ts = item.get("Date") - if not url or not ts: - continue - outpath = utc_filename(ts, url) - tasks_to_download.append((url, outpath)) + sem = asyncio.Semaphore(CONCURRENCY) + async with sem: + getting_urls = tqdm(total=len(memories), desc="Retrieving", unit="file") + # Prepare tasks + tasks_to_download = [] + + for item in memories: + url = item.get("Download Link") + ts = item.get("Date") + if not url or not ts: + continue + tasks_to_download.append((url, ts)) + getting_urls.update(1) + getting_urls.close() stats = { "mb": 0.0, @@ -112,19 +113,19 @@ async def main(): # tqdm based on number of files not yet done to_process = [ - (url, path) for (url, path) in tasks_to_download - if not path.exists() and str(path) not in stats["done"] + # if not path.exists() and str(path) not in stats["done"] + (url, timestamp) for (url, timestamp) in tasks_to_download ] progress = tqdm(total=len(to_process), desc="Downloading", unit="file") # Wrap each task so progress bar updates *after* the download completes - async def wrapped(url, path): - await download_one(session, sem, url, path, failures, stats) + async def wrapped(url, timestamp): + await download_one(session, sem, url, timestamp, failures, stats) progress.update(1) # Start tasks in batches - await asyncio.gather(*(wrapped(url, path) for url, path in to_process)) + await asyncio.gather(*(wrapped(url, timestamp) for url, timestamp in to_process)) progress.close() elapsed = time.time() - start_time @@ -141,7 +142,7 @@ async def main(): print("=" * 60) if failures: - print("\n⚠️ Failed downloads:") + print("\nFailed downloads:") for url, err in failures: print(f" - {url} ({err})") print()