renamed main

main.py

@@ -1,56 +1,37 @@
-import argparse
 import asyncio
+import aiohttp
 import json
-import os
 import pytz
-import re
 import time
-import subprocess
-from datetime import datetime
-from pathlib import Path
-
-import exif
 import httpx
-from pydantic import BaseModel, Field, field_validator
-from tqdm.asyncio import tqdm
+import os
+from pathlib import Path
+from datetime import datetime
+from tqdm import tqdm
+
+CONCURRENCY = 50
+RETRIES = 3
+OUTPUT_DIR = Path("./downloads")
+CHECKPOINT = Path("./temp/checkpoint.txt")
 
 
-class Memory(BaseModel):
-    date: datetime = Field(alias="Date")
-    download_link: str = Field(alias="Download Link")
-    location: str = Field(default="", alias="Location")
-    latitude: float | None = None
-    longitude: float | None = None
-
-    @field_validator("date", mode="before")
-    @classmethod
-    def parse_date(cls, v):
-        if isinstance(v, str):
-            return datetime.strptime(v, "%Y-%m-%d %H:%M:%S UTC")
-        return v
-
-    def model_post_init(self, __context):
-        if self.location and not self.latitude:
-            if match := re.search(r"([-\d.]+),\s*([-\d.]+)", self.location):
-                self.latitude = float(match.group(1))
-                self.longitude = float(match.group(2))
-
-    @property
-    def filename(self) -> str:
-        return self.date.strftime("%Y-%m-%d_%H-%M-%S")
+def load_checkpoint():
+    if CHECKPOINT.exists():
+        return set(CHECKPOINT.read_text().splitlines())
+    return set()
 
 
-class Stats(BaseModel):
-    downloaded: int = 0
-    skipped: int = 0
-    failed: int = 0
-    mb: float = 0
+def save_to_checkpoint(path: Path):
+    with CHECKPOINT.open("a") as f:
+        f.write(str(path) + "\n")
 
 
-def load_memories(json_path: Path) -> list[Memory]:
-    with open(json_path, "r", encoding="utf-8") as f:
-        data = json.load(f)
-    return [Memory(**item) for item in data["Saved Media"]]
+async def utc_filename(timestamp: str, url: str) -> Path:
+    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}")
+    return Path(out_path)
 
 
 async def get_cdn_url(download_link: str) -> str:
@@ -63,212 +44,97 @@ async def get_cdn_url(download_link: str) -> str:
         return response.text.strip()
 
 
-def add_exif_data(image_path: Path, memory: Memory):
-    try:
-        with open(image_path, "rb") as f:
-            img = exif.Image(f)
+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)
+                url = cdn_url
+                path = await utc_filename(timestamp, url)
+                async with session.get(url) as resp:
+                    resp.raise_for_status()
+                    data = await resp.read()
 
-        utc = pytz.utc
-        # Use the 'America/Chicago' zone which handles both CST (-6) and CDT (-5) automatically
-        central_timezone = pytz.timezone("America/Chicago")
-        utc_dt = memory.date.replace(tzinfo=utc)
-        local_dt = utc_dt.astimezone(central_timezone)
-        dt_str = local_dt.strftime("%Y:%m:%d %H:%M:%S")
-        # dt_str = memory.date.strftime("%Y:%m:%d %H:%M:%S")
-        img.datetime_original = dt_str
-        img.datetime_digitized = dt_str
-        img.datetime = dt_str
+                if path.exists():
+                    return
 
-        if memory.latitude is not None and memory.longitude is not None:
-            # Convert decimal degrees to degrees, minutes, seconds
-            def decimal_to_dms(decimal):
-                degrees = int(abs(decimal))
-                minutes_decimal = (abs(decimal) - degrees) * 60
-                minutes = int(minutes_decimal)
-                seconds = (minutes_decimal - minutes) * 60
-                return (degrees, minutes, seconds)
+                if str(path) in stats["done"]:
+                    return
 
-            lat_dms = decimal_to_dms(memory.latitude)
-            lon_dms = decimal_to_dms(memory.longitude)
+                path.parent.mkdir(parents=True, exist_ok=True)
+                path.write_bytes(data)
+                stats["mb"] += len(data) / (1024 * 1024)
+                save_to_checkpoint(path)
+                stats["done"].add(str(path))
 
-            img.gps_latitude = lat_dms
-            img.gps_latitude_ref = "N" if memory.latitude >= 0 else "S"
-            img.gps_longitude = lon_dms
-            img.gps_longitude_ref = "E" if memory.longitude >= 0 else "W"
-
-        with open(image_path, "wb") as f:
-            f.write(img.get_file())
-    except:
-        pass
-
-
-def add_metadata_mp4(video_path: Path, memory: Memory):
-    try:
-        dt_str = memory.date.strftime("%Y:%m:%d %H:%M:%S")
-
-        command = [
-            "exiftool",
-            "-P",  # Preserve file modification date
-            "-overwrite_original",
-            "-m",  # Allow minor warnings
-            # QuickTime Date/Time tags
-            f"-QuickTime:DateTimeOriginal={dt_str}",
-            f"-QuickTime:CreateDate={dt_str}",
-            f"-QuickTime:ModifyDate={dt_str}",
-            # QuickTime TimeZone is often required for the date to be recognized
-            f"-QuickTime:TimeZone=+00:00",
-            # # Final path argument
-            # str(video_path)
-        ]
-
-        # # 3. Add the target file path at the end
-        command.append(str(video_path))
-
-        # 4. Execute the command
-        # We use check=True to raise an error if ExifTool fails.
-        subprocess.run(
-            command,
-            check=True,
-            capture_output=True,  # Captures stdout/stderr
-            text=True,
-        )
-
-    except subprocess.CalledProcessError as e:
-        # Handle cases where ExifTool runs but returns an error (e.g., file corrupted)
-        print(f"\nError: ExifTool failed for {video_path.name}: {e.stderr.strip()}")
-    except FileNotFoundError:
-        # Handle case where the 'exiftool' executable cannot be found
-        print(
-            "\nError: ExifTool not found. Please install ExifTool and ensure it's in your system PATH."
-        )
-    except Exception as e:
-        # Catch all other exceptions
-        print(f"\nWarning: Failed to add metadata to {video_path.name}: {e}")
-        pass
-
-
-async def download_memory(
-    memory: Memory, output_dir: Path, add_exif: bool, semaphore: asyncio.Semaphore
-) -> tuple[bool, int]:
-    async with semaphore:
-        try:
-            cdn_url = await get_cdn_url(memory.download_link)
-            ext = Path(cdn_url.split("?")[0]).suffix or ".jpg"
-            output_path = output_dir / f"{memory.filename}{ext}"
-
-            async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
-                response = await client.get(cdn_url)
-                response.raise_for_status()
-
-                output_path.write_bytes(response.content)
-
-                timestamp = memory.date.timestamp()
-                os.utime(output_path, (timestamp, timestamp))
-
-                # if add_exif and ext == ".jpg":
-                #     add_exif_data(output_path, memory)
-
-                # elif add_exif and ext == ".mp4":
-                #     add_metadata_mp4(output_path, memory)
-
-                return True, len(response.content)
-        except Exception as e:
-            print(f"\nError: {e}")
-            return False, 0
-
-
-async def download_all(
-    memories: list[Memory],
-    output_dir: Path,
-    max_concurrent: int,
-    add_exif: bool,
-    skip_existing: bool,
-):
-    output_dir.mkdir(parents=True, exist_ok=True)
-    semaphore = asyncio.Semaphore(max_concurrent)
-    stats = Stats()
-    start_time = time.time()
-
-    to_download = []
-    for memory in memories:
-        jpg_path = output_dir / f"{memory.filename}.jpg"
-        mp4_path = output_dir / f"{memory.filename}.mp4"
-        if skip_existing and (jpg_path.exists() or mp4_path.exists()):
-            stats.skipped += 1
-        else:
-            to_download.append(memory)
-
-    if not to_download:
-        print("All files already downloaded!")
-        return
-
-    progress_bar = tqdm(
-        total=len(to_download),
-        desc="Downloading",
-        unit="file",
-        disable=False,
-    )
-
-    async def process_and_update(memory):
-        success, bytes_downloaded = await download_memory(
-            memory, output_dir, add_exif, semaphore
-        )
-        if success:
-            stats.downloaded += 1
-        else:
-            stats.failed += 1
-        stats.mb += bytes_downloaded / 1024 / 1024
-
-        elapsed = time.time() - start_time
-        mb_per_sec = (stats.mb) / elapsed if elapsed > 0 else 0
-        progress_bar.set_postfix({"MB/s": f"{mb_per_sec:.2f}"}, refresh=False)
-        progress_bar.update(1)
-
-    await asyncio.gather(*[process_and_update(m) for m in to_download])
-
-    progress_bar.close()
-    elapsed = time.time() - start_time
-    mb_total = stats.mb
-    mb_per_sec = mb_total / elapsed if elapsed > 0 else 0
-    print(
-        f"\n{'='*50}\nDownloaded: {stats.downloaded} ({mb_total:.1f} MB @ {mb_per_sec:.2f} MB/s) | Skipped: {stats.skipped} | Failed: {stats.failed}\n{'='*50}"
-    )
+                return
+            except Exception as e:
+                if attempt == RETRIES:
+                    failures.append((url, str(e)))
+                await asyncio.sleep(0.3 * attempt)
 
 
 async def main():
-    parser = argparse.ArgumentParser(
-        description="Download Snapchat memories from data export"
-    )
-    parser.add_argument(
-        "json_file",
-        nargs="?",
-        default="json/memories_history.json",
-        help="Path to memories_history.json",
-    )
-    parser.add_argument(
-        "-o", "--output", default="./downloads", help="Output directory"
-    )
-    parser.add_argument(
-        "-c", "--concurrent", type=int, default=40, help="Max concurrent downloads"
-    )
-    parser.add_argument("--no-exif", action="store_true", help="Disable EXIF metadata")
-    parser.add_argument(
-        "--no-skip-existing", action="store_true", help="Re-download existing files"
-    )
-    args = parser.parse_args()
+    with open("./resources/json/memories_history.json", "r") as f:
+        memories = json.load(f)["Saved Media"]
 
-    json_path = Path(args.json_file)
-    output_dir = Path(args.output)
+    sem = asyncio.Semaphore(CONCURRENCY)
+    async with sem:
+        getting_urls = tqdm(total=len(memories), desc="Retrieving", unit="file")
 
-    memories = load_memories(json_path)
+        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()
 
-    await download_all(
-        memories,
-        output_dir,
-        args.concurrent,
-        not args.no_exif,
-        not args.no_skip_existing,
-    )
+    stats = {
+        "mb": 0.0,
+        "done": load_checkpoint()
+    }
+
+    failures = []
+    sem = asyncio.Semaphore(CONCURRENCY)
+
+    async with aiohttp.ClientSession() as session:
+        start_time = time.time()
+
+        to_process = [
+            # 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")
+
+        async def wrapped(url, timestamp):
+            await download_one(session, sem, url, timestamp, failures, stats)
+            progress.update(1)
+
+        await asyncio.gather(*(wrapped(url, timestamp) for url, timestamp in to_process))
+        progress.close()
+
+        elapsed = time.time() - start_time
+
+    mb_total = stats["mb"]
+    speed = mb_total / elapsed if elapsed > 0 else 0
+
+    print("\n" + "=" * 60)
+    print(f"Downloaded: {len(to_process) - len(failures)} files")
+    print(f"Failed:     {len(failures)} files")
+    print(f"Data:       {mb_total:.2f} MB")
+    print(f"Speed:      {speed:.2f} MB/s")
+    print("=" * 60)
+
+    if failures:
+        print("\nFailed downloads:")
+        for url, err in failures:
+            print(f" - {url}   ({err})")
+        print()
 
 
 if __name__ == "__main__":

main3.py

@@ -1,141 +0,0 @@
-import asyncio
-import aiohttp
-import json
-import pytz
-import time
-import httpx
-import os
-from pathlib import Path
-from datetime import datetime
-from tqdm import tqdm
-
-CONCURRENCY = 50
-RETRIES = 3
-OUTPUT_DIR = Path("./downloads")
-CHECKPOINT = Path("./temp/checkpoint.txt")
-
-
-def load_checkpoint():
-    if CHECKPOINT.exists():
-        return set(CHECKPOINT.read_text().splitlines())
-    return set()
-
-
-def save_to_checkpoint(path: Path):
-    with CHECKPOINT.open("a") as f:
-        f.write(str(path) + "\n")
-
-
-async def utc_filename(timestamp: str, url: str) -> Path:
-    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}")
-    return Path(out_path)
-
-
-async def get_cdn_url(download_link: str) -> str:
-    async with httpx.AsyncClient(timeout=30.0) as client:
-        response = await client.post(
-            download_link,
-            headers={"Content-Type": "application/x-www-form-urlencoded"},
-        )
-        response.raise_for_status()
-        return response.text.strip()
-
-
-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)
-                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
-
-                if str(path) in stats["done"]:
-                    return
-
-                path.parent.mkdir(parents=True, exist_ok=True)
-                path.write_bytes(data)
-                stats["mb"] += len(data) / (1024 * 1024)
-                save_to_checkpoint(path)
-                stats["done"].add(str(path))
-
-                return
-            except Exception as e:
-                if attempt == RETRIES:
-                    failures.append((url, str(e)))
-                await asyncio.sleep(0.3 * attempt)
-
-
-async def main():
-    with open("./resources/json/memories_history.json", "r") as f:
-        memories = json.load(f)["Saved Media"]
-
-    sem = asyncio.Semaphore(CONCURRENCY)
-    async with sem:
-        getting_urls = tqdm(total=len(memories), desc="Retrieving", unit="file")
-
-        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,
-        "done": load_checkpoint()
-    }
-
-    failures = []
-    sem = asyncio.Semaphore(CONCURRENCY)
-
-    async with aiohttp.ClientSession() as session:
-        start_time = time.time()
-
-        to_process = [
-            # 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")
-
-        async def wrapped(url, timestamp):
-            await download_one(session, sem, url, timestamp, failures, stats)
-            progress.update(1)
-
-        await asyncio.gather(*(wrapped(url, timestamp) for url, timestamp in to_process))
-        progress.close()
-
-        elapsed = time.time() - start_time
-
-    mb_total = stats["mb"]
-    speed = mb_total / elapsed if elapsed > 0 else 0
-
-    print("\n" + "=" * 60)
-    print(f"Downloaded: {len(to_process) - len(failures)} files")
-    print(f"Failed:     {len(failures)} files")
-    print(f"Data:       {mb_total:.2f} MB")
-    print(f"Speed:      {speed:.2f} MB/s")
-    print("=" * 60)
-
-    if failures:
-        print("\nFailed downloads:")
-        for url, err in failures:
-            print(f" - {url}   ({err})")
-        print()
-
-
-if __name__ == "__main__":
-    asyncio.run(main())