TwitchDownloader/modules/utils.py

347 lines
11 KiB
Python
Raw Normal View History

"""
Utility functions and helpers for Twitch Archive.
"""
import os
import sys
import shutil
import pathlib
import subprocess
from typing import Optional
from colorama import Fore, Style
def detect_operating_system() -> str:
"""
Detect the current operating system.
Returns:
str: 'windows' or 'linux'
Raises:
SystemExit: If OS is not supported
"""
if sys.platform.startswith('win32'):
return 'windows'
elif sys.platform.startswith('linux'):
return 'linux'
else:
print(f'{Fore.RED}✗ ERROR: Unsupported operating system: {sys.platform}{Style.RESET_ALL}')
print(f'{Fore.YELLOW} This script only supports Windows and Linux{Style.RESET_ALL}')
sys.exit(1)
def get_bin_path() -> str:
"""Get the path to the bin directory containing external tools."""
return str(pathlib.Path(__file__).parent.parent.resolve() / "bin")
def get_env_value(*names: str, default: Optional[str] = None) -> Optional[str]:
"""Return the first non-empty environment variable from the provided names."""
for name in names:
value = os.getenv(name)
if value not in (None, ""):
return value
return default
def get_ffmpeg_executable(os_type: str) -> str:
"""
Get the platform-specific ffmpeg executable path.
Args:
os_type: Operating system type ('windows' or 'linux')
Returns:
str: Path to ffmpeg executable
"""
bin_path = get_bin_path()
if os_type == 'windows':
return os.path.join(bin_path, 'ffmpeg.exe')
system_ffmpeg = shutil.which('ffmpeg')
if system_ffmpeg:
return system_ffmpeg
return os.path.join(bin_path, 'ffmpeg')
def get_twitch_downloader_executable(os_type: str) -> str:
"""
Get the platform-specific TwitchDownloaderCLI executable path.
Args:
os_type: Operating system type ('windows' or 'linux')
Returns:
str: Path to TwitchDownloaderCLI executable
"""
bin_path = get_bin_path()
if os_type == 'windows':
return os.path.join(bin_path, 'TwitchDownloaderCLI.exe')
system_twitch_downloader = shutil.which('TwitchDownloaderCLI')
if system_twitch_downloader:
return system_twitch_downloader
return os.path.join(bin_path, 'TwitchDownloaderCLI')
def get_unique_filename(filepath: str) -> str:
"""
Generate a unique filename by appending a counter if file already exists.
Args:
filepath: The desired file path
Returns:
str: A unique file path (original or with _N suffix)
Example:
If 'video.mp4' exists, returns 'video_1.mp4'
If 'video_1.mp4' also exists, returns 'video_2.mp4'
"""
if not os.path.exists(filepath):
return filepath
# Split into components
directory = os.path.dirname(filepath)
filename = os.path.basename(filepath)
name, ext = os.path.splitext(filename)
# Find next available counter
counter = 1
while True:
new_filepath = os.path.join(directory, f"{name}_{counter}{ext}")
if not os.path.exists(new_filepath):
return new_filepath
counter += 1
def get_video_duration(video_path: str, ffmpeg_path: str) -> Optional[float]:
"""
Get the duration of a video file in seconds using FFmpeg.
Args:
video_path: Path to the video file
ffmpeg_path: Path to FFmpeg executable
Returns:
float: Duration in seconds, or None if failed
"""
if not os.path.exists(video_path):
return None
try:
# Use ffprobe (comes with ffmpeg) to get duration
ffprobe_path = ffmpeg_path.replace('ffmpeg', 'ffprobe')
cmd = [
ffprobe_path,
'-v', 'error',
'-show_entries', 'format=duration',
'-of', 'default=noprint_wrappers=1:nokey=1',
video_path
]
result = subprocess.run(
cmd,
capture_output=True,
text=True,
timeout=30
)
if result.returncode == 0 and result.stdout.strip():
return float(result.stdout.strip())
except Exception:
pass
return None
def verify_streamlink() -> bool:
"""
Verify that streamlink is available.
Returns:
bool: True if streamlink is available, False otherwise
"""
try:
result = subprocess.run(['streamlink', '--version'],
capture_output=True,
text=True,
timeout=5)
if result.returncode == 0:
version = result.stdout.strip().split()[1] if len(result.stdout.split()) > 1 else 'unknown'
print(f'{Fore.GREEN}✓ Streamlink v{version} found{Style.RESET_ALL}')
return True
else:
raise FileNotFoundError()
except (FileNotFoundError, subprocess.TimeoutExpired, IndexError):
print(f'{Fore.RED}✗ ERROR: Streamlink not found{Style.RESET_ALL}')
print(f'{Fore.CYAN} → Install streamlink: pip install streamlink{Style.RESET_ALL}')
print(f'{Fore.CYAN} → Or download from: https://streamlink.github.io/{Style.RESET_ALL}')
return False
def verify_rclone() -> bool:
"""Verify that rclone is available on PATH."""
try:
result = subprocess.run(['rclone', 'version'],
capture_output=True,
text=True,
timeout=5)
if result.returncode == 0:
version_line = result.stdout.strip().splitlines()[0] if result.stdout.strip() else 'unknown'
print(f'{Fore.GREEN}✓ Rclone found ({version_line}){Style.RESET_ALL}')
return True
raise FileNotFoundError()
except (FileNotFoundError, subprocess.TimeoutExpired, IndexError):
print(f'{Fore.RED}✗ ERROR: rclone not found{Style.RESET_ALL}')
print(f'{Fore.CYAN} → Install rclone and ensure it is on PATH{Style.RESET_ALL}')
return False
def verify_ffmpeg(os_type: str) -> bool:
"""
Verify that ffmpeg is available.
Args:
os_type: Operating system type ('windows' or 'linux')
Returns:
bool: True if ffmpeg is available, False otherwise
"""
try:
ffmpeg_path = get_ffmpeg_executable(os_type)
if os.path.exists(ffmpeg_path):
print(f'{Fore.GREEN}✓ FFmpeg found at {ffmpeg_path}{Style.RESET_ALL}')
return True
else:
print(f'{Fore.YELLOW}⚠ Warning: FFmpeg not found at {ffmpeg_path}{Style.RESET_ALL}')
print(f'{Fore.YELLOW} → Download FFmpeg and place it in the bin/ folder{Style.RESET_ALL}')
return False
except Exception as e:
print(f'{Fore.YELLOW}⚠ Warning: Could not verify FFmpeg: {e}{Style.RESET_ALL}')
return False
def verify_twitch_downloader(os_type: str) -> bool:
"""
Verify that TwitchDownloaderCLI is available.
Args:
os_type: Operating system type ('windows' or 'linux')
Returns:
bool: True if TwitchDownloaderCLI is available, False otherwise
"""
try:
downloader_path = get_twitch_downloader_executable(os_type)
if os.path.exists(downloader_path):
print(f'{Fore.GREEN}✓ TwitchDownloaderCLI found{Style.RESET_ALL}')
return True
else:
print(f'{Fore.YELLOW}⚠ Warning: TwitchDownloaderCLI not found at {downloader_path}{Style.RESET_ALL}')
print(f'{Fore.YELLOW} → Download from: https://github.com/lay295/TwitchDownloader/releases{Style.RESET_ALL}')
return False
except Exception as e:
print(f'{Fore.YELLOW}⚠ Warning: Could not verify TwitchDownloaderCLI: {e}{Style.RESET_ALL}')
return False
def detect_hardware_acceleration(hwaccel_config: str, os_type: str) -> Optional[str]:
"""
Detect available hardware acceleration based on config and system.
Args:
hwaccel_config: Hardware acceleration configuration ('auto', 'nvenc', 'qsv', 'amf', 'vaapi', 'none')
os_type: Operating system type ('windows' or 'linux')
Returns:
str: Hardware acceleration type or None
"""
# If user explicitly set to 'none', disable hardware acceleration
if hwaccel_config == 'none':
return 'none'
# If user specified a particular type, use it
if hwaccel_config in ['nvenc', 'qsv', 'amf', 'vaapi']:
return hwaccel_config
# Auto-detect: choose only hardware we can reasonably prove is present.
if hwaccel_config == 'auto':
if is_nvidia_runtime_available():
return 'nvenc'
if is_vaapi_runtime_available():
return 'vaapi'
return 'none'
return None
def is_nvidia_runtime_available() -> bool:
"""Return True when the current runtime appears to expose an NVIDIA GPU."""
visible_devices = os.getenv('NVIDIA_VISIBLE_DEVICES', '').strip().lower()
if visible_devices in {'void', 'none'}:
return False
if visible_devices and visible_devices != 'all':
return True
if shutil.which('nvidia-smi'):
return True
return any(
os.path.exists(device_path)
for device_path in ('/dev/nvidiactl', '/dev/nvidia0', '/dev/nvidia-modeset')
)
def is_vaapi_runtime_available() -> bool:
"""Return True when Linux VAAPI render nodes are present."""
return any(
os.path.exists(device_path)
for device_path in ('/dev/dri/renderD128', '/dev/dri/card0')
)
def resolve_hwaccel_type(hwaccel_type: Optional[str], os_type: str) -> Optional[str]:
"""Return a safe hardware acceleration choice for the current runtime."""
if hwaccel_type in (None, 'none'):
return 'none'
if hwaccel_type == 'nvenc':
return 'nvenc' if is_nvidia_runtime_available() else 'none'
if hwaccel_type == 'vaapi':
return 'vaapi' if is_vaapi_runtime_available() else 'none'
# Leave explicit QSV/AMF unchanged for non-container users; container auto-detect no longer picks them blindly.
return hwaccel_type
def get_hwaccel_encoder(hwaccel_type: str) -> str:
"""
Get the appropriate hardware-accelerated encoder for the given acceleration type.
Args:
hwaccel_type: Type of hardware acceleration ('nvenc', 'qsv', 'amf', 'vaapi', 'auto', 'none')
Returns:
str: FFmpeg encoder name (e.g., 'h264_nvenc', 'libx264')
"""
encoder_map = {
'nvenc': 'h264_nvenc', # NVIDIA
'qsv': 'h264_qsv', # Intel Quick Sync
'amf': 'h264_amf', # AMD
'vaapi': 'h264_vaapi', # Linux VA-API
}
if hwaccel_type in encoder_map:
return encoder_map[hwaccel_type]
elif hwaccel_type == 'auto':
# Try NVENC first (most common), fall back to libx264
# In real usage, auto will attempt to use what's available
return 'h264_nvenc'
else:
return 'libx264' # Software encoding fallback