removed old main
This commit is contained in:
parent
832472d48d
commit
11a8c4e0a0
275
main.py
275
main.py
@ -1,275 +0,0 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
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
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
class Stats(BaseModel):
|
||||
downloaded: int = 0
|
||||
skipped: int = 0
|
||||
failed: int = 0
|
||||
mb: float = 0
|
||||
|
||||
|
||||
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 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()
|
||||
|
||||
|
||||
def add_exif_data(image_path: Path, memory: Memory):
|
||||
try:
|
||||
with open(image_path, "rb") as f:
|
||||
img = exif.Image(f)
|
||||
|
||||
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 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)
|
||||
|
||||
lat_dms = decimal_to_dms(memory.latitude)
|
||||
lon_dms = decimal_to_dms(memory.longitude)
|
||||
|
||||
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}"
|
||||
)
|
||||
|
||||
|
||||
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()
|
||||
|
||||
json_path = Path(args.json_file)
|
||||
output_dir = Path(args.output)
|
||||
|
||||
memories = load_memories(json_path)
|
||||
|
||||
await download_all(
|
||||
memories,
|
||||
output_dir,
|
||||
args.concurrent,
|
||||
not args.no_exif,
|
||||
not args.no_skip_existing,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Loading…
Reference in New Issue
Block a user