Embrace pathlib

This commit is contained in:
Ivan Habunek 2024-08-28 10:59:23 +02:00
parent 29d3f7fec2
commit 1658aba124
No known key found for this signature in database
GPG Key ID: 01DB3DD0D824504C
5 changed files with 32 additions and 29 deletions

View File

@ -1,6 +1,7 @@
import re import re
import sys import sys
from os import path from os import path
from pathlib import Path
from typing import Callable, Generator, Optional from typing import Callable, Generator, Optional
import click import click
@ -69,9 +70,9 @@ def _target_filename(clip: Clip):
def _download_clips(generator: Generator[Clip, None, None]): def _download_clips(generator: Generator[Clip, None, None]):
for clip in generator: for clip in generator:
target = _target_filename(clip) target = Path(_target_filename(clip))
if path.exists(target): if target.exists():
click.echo(f"Already downloaded: {green(target)}") click.echo(f"Already downloaded: {green(target)}")
else: else:
url = get_clip_authenticated_url(clip["slug"], "source") url = get_clip_authenticated_url(clip["slug"], "source")

View File

@ -1,11 +1,9 @@
import asyncio import asyncio
import os
import platform import platform
import re import re
import shutil import shutil
import subprocess import subprocess
import tempfile import tempfile
from os import path
from pathlib import Path from pathlib import Path
from typing import List, Optional from typing import List, Optional
from urllib.parse import urlencode, urlparse from urllib.parse import urlencode, urlparse
@ -51,14 +49,14 @@ def download_one(video: str, args: DownloadOptions):
raise ConsoleError(f"Invalid input: {video}") raise ConsoleError(f"Invalid input: {video}")
def _join_vods(playlist_path: str, target: str, overwrite: bool, video: Video): def _join_vods(playlist_path: Path, target: str, overwrite: bool, video: Video):
description = video["description"] or "" description = video["description"] or ""
description = description.strip() description = description.strip()
command = [ command: List[str] = [
"ffmpeg", "ffmpeg",
"-i", "-i",
playlist_path, str(playlist_path),
"-c", "-c",
"copy", "copy",
"-metadata", "-metadata",
@ -84,9 +82,9 @@ def _join_vods(playlist_path: str, target: str, overwrite: bool, video: Video):
raise ConsoleError("Joining files failed") raise ConsoleError("Joining files failed")
def _concat_vods(vod_paths: List[str], target: str): def _concat_vods(vod_paths: List[Path], target: str):
tool = "type" if platform.system() == "Windows" else "cat" tool = "type" if platform.system() == "Windows" else "cat"
command = [tool] + vod_paths command = [tool] + [str(p) for p in vod_paths]
with open(target, "wb") as target_file: with open(target, "wb") as target_file:
result = subprocess.run(command, stdout=target_file) result = subprocess.run(command, stdout=target_file)
@ -94,12 +92,12 @@ def _concat_vods(vod_paths: List[str], target: str):
raise ConsoleError(f"Joining files failed: {result.stderr}") raise ConsoleError(f"Joining files failed: {result.stderr}")
def _crete_temp_dir(base_uri: str) -> str: def _crete_temp_dir(base_uri: str) -> Path:
"""Create a temp dir to store downloads if it doesn't exist.""" """Create a temp dir to store downloads if it doesn't exist."""
path = urlparse(base_uri).path.lstrip("/") path = urlparse(base_uri).path.lstrip("/")
temp_dir = Path(tempfile.gettempdir(), "twitch-dl", path) temp_dir = Path(tempfile.gettempdir(), "twitch-dl", path)
temp_dir.mkdir(parents=True, exist_ok=True) temp_dir.mkdir(parents=True, exist_ok=True)
return str(temp_dir) return temp_dir
def _get_clip_url(access_token: ClipAccessToken, quality: Optional[str]) -> str: def _get_clip_url(access_token: ClipAccessToken, quality: Optional[str]) -> str:
@ -162,10 +160,10 @@ def _download_clip(slug: str, args: DownloadOptions) -> None:
duration = utils.format_duration(clip["durationSeconds"]) duration = utils.format_duration(clip["durationSeconds"])
click.echo(f"Found: {green(title)} by {yellow(user)}, playing {blue(game)} ({duration})") click.echo(f"Found: {green(title)} by {yellow(user)}, playing {blue(game)} ({duration})")
target = clip_filename(clip, args.output) target = Path(clip_filename(clip, args.output))
click.echo(f"Target: {blue(target)}") click.echo(f"Target: {blue(target)}")
if not args.overwrite and path.exists(target): if not args.overwrite and target.exists():
response = click.prompt("File exists. Overwrite? [Y/n]", default="Y", show_default=False) response = click.prompt("File exists. Overwrite? [Y/n]", default="Y", show_default=False)
if response.lower().strip() != "y": if response.lower().strip() != "y":
raise click.Abort() raise click.Abort()
@ -194,10 +192,10 @@ def _download_video(video_id: str, args: DownloadOptions) -> None:
click.echo(f"Found: {blue(video['title'])} by {yellow(video['creator']['displayName'])}") click.echo(f"Found: {blue(video['title'])} by {yellow(video['creator']['displayName'])}")
target = video_filename(video, args.format, args.output) target = Path(video_filename(video, args.format, args.output))
click.echo(f"Output: {blue(target)}") click.echo(f"Output: {blue(target)}")
if not args.overwrite and path.exists(target): if not args.overwrite and target.exists():
response = click.prompt("File exists. Overwrite? [Y/n]", default="Y", show_default=False) response = click.prompt("File exists. Overwrite? [Y/n]", default="Y", show_default=False)
if response.lower().strip() != "y": if response.lower().strip() != "y":
raise click.Abort() raise click.Abort()
@ -227,19 +225,19 @@ def _download_video(video_id: str, args: DownloadOptions) -> None:
target_dir = _crete_temp_dir(base_uri) target_dir = _crete_temp_dir(base_uri)
# Save playlists for debugging purposes # Save playlists for debugging purposes
with open(path.join(target_dir, "playlists.m3u8"), "w") as f: with open(target_dir / "playlists.m3u8", "w") as f:
f.write(playlists_text) f.write(playlists_text)
with open(path.join(target_dir, "playlist.m3u8"), "w") as f: with open(target_dir / "playlist.m3u8", "w") as f:
f.write(vods_text) f.write(vods_text)
click.echo(f"\nDownloading {len(vods)} VODs using {args.max_workers} workers to {target_dir}") click.echo(f"\nDownloading {len(vods)} VODs using {args.max_workers} workers to {target_dir}")
sources = [base_uri + vod.path for vod in vods] sources = [base_uri + vod.path for vod in vods]
targets = [os.path.join(target_dir, f"{vod.index:05d}.ts") for vod in vods] targets = [target_dir / f"{vod.index:05d}.ts" for vod in vods]
asyncio.run(download_all(sources, targets, args.max_workers, rate_limit=args.rate_limit)) asyncio.run(download_all(sources, targets, args.max_workers, rate_limit=args.rate_limit))
join_playlist = make_join_playlist(vods_m3u8, vods, targets) join_playlist = make_join_playlist(vods_m3u8, vods, targets)
join_playlist_path = path.join(target_dir, "playlist_downloaded.m3u8") join_playlist_path = target_dir / "playlist_downloaded.m3u8"
join_playlist.dump(join_playlist_path) # type: ignore join_playlist.dump(join_playlist_path) # type: ignore
click.echo() click.echo()

View File

@ -1,4 +1,6 @@
import os import os
from pathlib import Path
from typing import Tuple
import httpx import httpx
@ -9,8 +11,8 @@ CONNECT_TIMEOUT = 5
RETRY_COUNT = 5 RETRY_COUNT = 5
def _download(url: str, path: str): def _download(url: str, path: Path):
tmp_path = path + ".tmp" tmp_path = Path(str(path) + ".tmp")
size = 0 size = 0
with httpx.stream("GET", url, timeout=CONNECT_TIMEOUT, follow_redirects=True) as response: with httpx.stream("GET", url, timeout=CONNECT_TIMEOUT, follow_redirects=True) as response:
response.raise_for_status() response.raise_for_status()
@ -23,8 +25,8 @@ def _download(url: str, path: str):
return size return size
def download_file(url: str, path: str, retries: int = RETRY_COUNT): def download_file(url: str, path: Path, retries: int = RETRY_COUNT) -> Tuple[int, bool]:
if os.path.exists(path): if path.exists():
from_disk = True from_disk = True
return os.path.getsize(path), from_disk return os.path.getsize(path), from_disk

View File

@ -3,6 +3,7 @@ import logging
import os import os
import time import time
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pathlib import Path
from typing import List, Optional from typing import List, Optional
import httpx import httpx
@ -71,7 +72,7 @@ async def download(
client: httpx.AsyncClient, client: httpx.AsyncClient,
task_id: int, task_id: int,
source: str, source: str,
target: str, target: Path,
progress: Progress, progress: Progress,
token_bucket: TokenBucket, token_bucket: TokenBucket,
): ):
@ -96,12 +97,12 @@ async def download_with_retries(
semaphore: asyncio.Semaphore, semaphore: asyncio.Semaphore,
task_id: int, task_id: int,
source: str, source: str,
target: str, target: Path,
progress: Progress, progress: Progress,
token_bucket: TokenBucket, token_bucket: TokenBucket,
): ):
async with semaphore: async with semaphore:
if os.path.exists(target): if target.exists():
size = os.path.getsize(target) size = os.path.getsize(target)
progress.already_downloaded(task_id, size) progress.already_downloaded(task_id, size)
return return
@ -120,7 +121,7 @@ async def download_with_retries(
async def download_all( async def download_all(
sources: List[str], sources: List[str],
targets: List[str], targets: List[Path],
workers: int, workers: int,
*, *,
rate_limit: Optional[int] = None, rate_limit: Optional[int] = None,

View File

@ -3,6 +3,7 @@ Parse and manipulate m3u8 playlists.
""" """
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path
from typing import Generator, List, Optional, OrderedDict from typing import Generator, List, Optional, OrderedDict
import click import click
@ -81,7 +82,7 @@ def enumerate_vods(
def make_join_playlist( def make_join_playlist(
playlist: m3u8.M3U8, playlist: m3u8.M3U8,
vods: List[Vod], vods: List[Vod],
targets: List[str], targets: List[Path],
) -> m3u8.Playlist: ) -> m3u8.Playlist:
""" """
Make a modified playlist which references downloaded VODs Make a modified playlist which references downloaded VODs
@ -93,7 +94,7 @@ def make_join_playlist(
playlist.segments.clear() playlist.segments.clear()
for segment in org_segments: for segment in org_segments:
if segment.uri in path_map: if segment.uri in path_map:
segment.uri = path_map[segment.uri] segment.uri = str(path_map[segment.uri])
playlist.segments.append(segment) playlist.segments.append(segment)
return playlist return playlist