This commit is contained in:
pacnpal
2024-11-17 01:03:45 +00:00
parent d2576df988
commit 2557978cf3
6 changed files with 143 additions and 451 deletions

View File

@@ -40,6 +40,7 @@ from .reactions import (
update_progress_reaction, update_progress_reaction,
update_download_progress_reaction update_download_progress_reaction
) )
from ..utils import progress_tracker
# Export public classes and constants # Export public classes and constants
__all__ = [ __all__ = [
@@ -207,6 +208,3 @@ def clear_caches(message_id: Optional[int] = None) -> None:
""" """
url_extractor.clear_cache(message_id) url_extractor.clear_cache(message_id)
message_validator.clear_cache(message_id) message_validator.clear_cache(message_id)
# Initialize shared progress tracker instance
progress_tracker = ProgressTracker()

View File

@@ -10,7 +10,7 @@ from discord.ext import commands
from .message_handler import MessageHandler from .message_handler import MessageHandler
from .queue_handler import QueueHandler from .queue_handler import QueueHandler
from ..utils.progress_tracker import ProgressTracker from ..utils import progress_tracker
from .status_display import StatusDisplay from .status_display import StatusDisplay
from .cleanup_manager import CleanupManager, CleanupStrategy from .cleanup_manager import CleanupManager, CleanupStrategy
from .constants import REACTIONS from .constants import REACTIONS
@@ -196,7 +196,7 @@ class HealthMonitor:
self.health_status.update({ self.health_status.update({
"queue_handler": self.processor.queue_handler.is_healthy(), "queue_handler": self.processor.queue_handler.is_healthy(),
"message_handler": self.processor.message_handler.is_healthy(), "message_handler": self.processor.message_handler.is_healthy(),
"progress_tracker": self.processor.progress_tracker.is_healthy() "progress_tracker": progress_tracker.is_healthy()
}) })
# Check operation health # Check operation health
@@ -248,7 +248,6 @@ class VideoProcessor:
# Initialize handlers # Initialize handlers
self.queue_handler = QueueHandler(bot, config_manager, components) self.queue_handler = QueueHandler(bot, config_manager, components)
self.message_handler = MessageHandler(bot, config_manager, queue_manager) self.message_handler = MessageHandler(bot, config_manager, queue_manager)
self.progress_tracker = ProgressTracker()
self.cleanup_manager = CleanupManager( self.cleanup_manager = CleanupManager(
self.queue_handler, self.queue_handler,
ffmpeg_mgr, ffmpeg_mgr,

View File

@@ -1,88 +0,0 @@
"""Progress tracking for video downloads and compression"""
from typing import Dict, Any
from datetime import datetime
class ProgressTracker:
"""Tracks progress of video downloads and compression operations"""
def __init__(self):
self._download_progress: Dict[str, Dict[str, Any]] = {}
self._compression_progress: Dict[str, Dict[str, Any]] = {}
def update_download_progress(self, url: str, progress_data: Dict[str, Any]) -> None:
"""Update download progress for a specific URL"""
if url not in self._download_progress:
self._download_progress[url] = {
'active': True,
'start_time': datetime.utcnow().isoformat(),
'retries': 0
}
self._download_progress[url].update(progress_data)
def complete_download(self, url: str) -> None:
"""Mark a download as complete"""
if url in self._download_progress:
self._download_progress[url]['active'] = False
self._download_progress[url]['completed_time'] = datetime.utcnow().isoformat()
def increment_download_retries(self, url: str) -> None:
"""Increment retry count for a download"""
if url in self._download_progress:
self._download_progress[url]['retries'] = self._download_progress[url].get('retries', 0) + 1
def update_compression_progress(self, file_id: str, progress_data: Dict[str, Any]) -> None:
"""Update compression progress for a specific file"""
if file_id not in self._compression_progress:
self._compression_progress[file_id] = {
'active': True,
'start_time': datetime.utcnow().isoformat()
}
self._compression_progress[file_id].update(progress_data)
def complete_compression(self, file_id: str) -> None:
"""Mark a compression operation as complete"""
if file_id in self._compression_progress:
self._compression_progress[file_id]['active'] = False
self._compression_progress[file_id]['completed_time'] = datetime.utcnow().isoformat()
def get_download_progress(self, url: str = None) -> Dict[str, Any]:
"""Get download progress for a specific URL or all downloads"""
if url:
return self._download_progress.get(url, {})
return self._download_progress
def get_compression_progress(self, file_id: str = None) -> Dict[str, Any]:
"""Get compression progress for a specific file or all compressions"""
if file_id:
return self._compression_progress.get(file_id, {})
return self._compression_progress
def clear_completed(self) -> None:
"""Clear completed operations from tracking"""
# Clear completed downloads
self._download_progress = {
url: data for url, data in self._download_progress.items()
if data.get('active', False)
}
# Clear completed compressions
self._compression_progress = {
file_id: data for file_id, data in self._compression_progress.items()
if data.get('active', False)
}
def get_active_operations(self) -> Dict[str, Dict[str, Any]]:
"""Get all active operations"""
return {
'downloads': {
url: data for url, data in self._download_progress.items()
if data.get('active', False)
},
'compressions': {
file_id: data for file_id, data in self._compression_progress.items()
if data.get('active', False)
}
}

View File

@@ -8,7 +8,7 @@ from typing import Optional, Dict, Any, List, Tuple, Set, TypedDict, ClassVar, C
from datetime import datetime from datetime import datetime
import discord import discord
from ..utils.progress_tracker import ProgressTracker from ..utils import progress_tracker
from ..database.video_archive_db import VideoArchiveDB from ..database.video_archive_db import VideoArchiveDB
from ..utils.download_manager import DownloadManager from ..utils.download_manager import DownloadManager
from ..utils.message_manager import MessageManager from ..utils.message_manager import MessageManager
@@ -57,7 +57,6 @@ class QueueHandler:
self._unloading = False self._unloading = False
self._active_downloads: Dict[str, asyncio.Task] = {} self._active_downloads: Dict[str, asyncio.Task] = {}
self._active_downloads_lock = asyncio.Lock() self._active_downloads_lock = asyncio.Lock()
self.progress_tracker = ProgressTracker()
self._stats: QueueStats = { self._stats: QueueStats = {
"active_downloads": 0, "active_downloads": 0,
"processing_items": 0, "processing_items": 0,
@@ -105,12 +104,16 @@ class QueueHandler:
message_manager = components.get("message_manager") message_manager = components.get("message_manager")
if not downloader or not message_manager: if not downloader or not message_manager:
raise QueueHandlerError(f"Missing required components for guild {item.guild_id}") raise QueueHandlerError(
f"Missing required components for guild {item.guild_id}"
)
# Get original message and update reactions # Get original message and update reactions
original_message = await self._get_original_message(item) original_message = await self._get_original_message(item)
if original_message: if original_message:
await self._update_message_reactions(original_message, QueueItemStatus.PROCESSING) await self._update_message_reactions(
original_message, QueueItemStatus.PROCESSING
)
# Download and archive video # Download and archive video
file_path = await self._process_video_file( file_path = await self._process_video_file(
@@ -121,7 +124,9 @@ class QueueHandler:
self._update_stats(True, start_time) self._update_stats(True, start_time)
item.finish_processing(True) item.finish_processing(True)
if original_message: if original_message:
await self._update_message_reactions(original_message, QueueItemStatus.COMPLETED) await self._update_message_reactions(
original_message, QueueItemStatus.COMPLETED
)
return True, None return True, None
except QueueHandlerError as e: except QueueHandlerError as e:
@@ -143,7 +148,9 @@ class QueueHandler:
if self.db.is_url_archived(item.url): if self.db.is_url_archived(item.url):
logger.info(f"Video already archived: {item.url}") logger.info(f"Video already archived: {item.url}")
if original_message := await self._get_original_message(item): if original_message := await self._get_original_message(item):
await self._update_message_reactions(original_message, QueueItemStatus.COMPLETED) await self._update_message_reactions(
original_message, QueueItemStatus.COMPLETED
)
archived_info = self.db.get_archived_video(item.url) archived_info = self.db.get_archived_video(item.url)
if archived_info: if archived_info:
await original_message.reply( await original_message.reply(
@@ -153,10 +160,7 @@ class QueueHandler:
return True return True
return False return False
async def _get_components( async def _get_components(self, guild_id: int) -> Dict[str, Any]:
self,
guild_id: int
) -> Dict[str, Any]:
"""Get required components for processing""" """Get required components for processing"""
if guild_id not in self.components: if guild_id not in self.components:
raise QueueHandlerError(f"No components found for guild {guild_id}") raise QueueHandlerError(f"No components found for guild {guild_id}")
@@ -167,7 +171,7 @@ class QueueHandler:
downloader: DownloadManager, downloader: DownloadManager,
message_manager: MessageManager, message_manager: MessageManager,
item: QueueItem, item: QueueItem,
original_message: Optional[discord.Message] original_message: Optional[discord.Message],
) -> Optional[str]: ) -> Optional[str]:
"""Download and process video file""" """Download and process video file"""
# Create progress callback # Create progress callback
@@ -182,11 +186,7 @@ class QueueHandler:
# Archive video # Archive video
success, error = await self._archive_video( success, error = await self._archive_video(
item.guild_id, item.guild_id, original_message, message_manager, item.url, file_path
original_message,
message_manager,
item.url,
file_path
) )
if not success: if not success:
raise QueueHandlerError(f"Failed to archive video: {error}") raise QueueHandlerError(f"Failed to archive video: {error}")
@@ -194,16 +194,15 @@ class QueueHandler:
return file_path return file_path
def _handle_processing_error( def _handle_processing_error(
self, self, item: QueueItem, message: Optional[discord.Message], error: str
item: QueueItem,
message: Optional[discord.Message],
error: str
) -> None: ) -> None:
"""Handle processing error""" """Handle processing error"""
self._update_stats(False, datetime.utcnow()) self._update_stats(False, datetime.utcnow())
item.finish_processing(False, error) item.finish_processing(False, error)
if message: if message:
asyncio.create_task(self._update_message_reactions(message, QueueItemStatus.FAILED)) asyncio.create_task(
self._update_message_reactions(message, QueueItemStatus.FAILED)
)
def _update_stats(self, success: bool, start_time: datetime) -> None: def _update_stats(self, success: bool, start_time: datetime) -> None:
"""Update queue statistics""" """Update queue statistics"""
@@ -218,14 +217,14 @@ class QueueHandler:
total_items = self._stats["completed_items"] + self._stats["failed_items"] total_items = self._stats["completed_items"] + self._stats["failed_items"]
if total_items > 0: if total_items > 0:
current_total = self._stats["average_processing_time"] * (total_items - 1) current_total = self._stats["average_processing_time"] * (total_items - 1)
self._stats["average_processing_time"] = (current_total + processing_time) / total_items self._stats["average_processing_time"] = (
current_total + processing_time
) / total_items
self._stats["last_processed"] = datetime.utcnow().isoformat() self._stats["last_processed"] = datetime.utcnow().isoformat()
async def _update_message_reactions( async def _update_message_reactions(
self, self, message: discord.Message, status: QueueItemStatus
message: discord.Message,
status: QueueItemStatus
) -> None: ) -> None:
"""Update message reactions based on status""" """Update message reactions based on status"""
try: try:
@@ -234,7 +233,7 @@ class QueueHandler:
REACTIONS["queued"], REACTIONS["queued"],
REACTIONS["processing"], REACTIONS["processing"],
REACTIONS["success"], REACTIONS["success"],
REACTIONS["error"] REACTIONS["error"],
]: ]:
try: try:
await message.remove_reaction(reaction, self.bot.user) await message.remove_reaction(reaction, self.bot.user)
@@ -265,7 +264,7 @@ class QueueHandler:
original_message: Optional[discord.Message], original_message: Optional[discord.Message],
message_manager: MessageManager, message_manager: MessageManager,
url: str, url: str,
file_path: str file_path: str,
) -> Tuple[bool, Optional[str]]: ) -> Tuple[bool, Optional[str]]:
""" """
Archive downloaded video. Archive downloaded video.
@@ -308,19 +307,14 @@ class QueueHandler:
raise QueueHandlerError("Processed file not found") raise QueueHandlerError("Processed file not found")
archive_message = await archive_channel.send( archive_message = await archive_channel.send(
content=message, content=message, file=discord.File(file_path)
file=discord.File(file_path)
) )
# Store in database if available # Store in database if available
if self.db and archive_message.attachments: if self.db and archive_message.attachments:
discord_url = archive_message.attachments[0].url discord_url = archive_message.attachments[0].url
self.db.add_archived_video( self.db.add_archived_video(
url, url, discord_url, archive_message.id, archive_channel.id, guild_id
discord_url,
archive_message.id,
archive_channel.id,
guild_id
) )
logger.info(f"Added video to archive database: {url} -> {discord_url}") logger.info(f"Added video to archive database: {url} -> {discord_url}")
@@ -333,10 +327,7 @@ class QueueHandler:
logger.error(f"Failed to archive video: {str(e)}") logger.error(f"Failed to archive video: {str(e)}")
raise QueueHandlerError(f"Failed to archive video: {str(e)}") raise QueueHandlerError(f"Failed to archive video: {str(e)}")
async def _get_original_message( async def _get_original_message(self, item: QueueItem) -> Optional[discord.Message]:
self,
item: QueueItem
) -> Optional[discord.Message]:
""" """
Retrieve the original message. Retrieve the original message.
@@ -358,9 +349,7 @@ class QueueHandler:
return None return None
def _create_progress_callback( def _create_progress_callback(
self, self, message: Optional[discord.Message], url: str
message: Optional[discord.Message],
url: str
) -> Callable[[float], None]: ) -> Callable[[float], None]:
""" """
Create progress callback function for download tracking. Create progress callback function for download tracking.
@@ -372,34 +361,40 @@ class QueueHandler:
Returns: Returns:
Callback function for progress updates Callback function for progress updates
""" """
def progress_callback(progress: float) -> None: def progress_callback(progress: float) -> None:
if message: if message:
try: try:
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
if not loop.is_running(): if not loop.is_running():
logger.warning("Event loop is not running, skipping progress update") logger.warning(
"Event loop is not running, skipping progress update"
)
return return
# Update progress tracking # Update progress tracking
self.progress_tracker.update_download_progress(url, { progress_tracker.update_download_progress(
'percent': progress, url,
'last_update': datetime.utcnow().isoformat() {
}) "percent": progress,
"last_update": datetime.utcnow().isoformat(),
},
)
# Create task to update reaction # Create task to update reaction
asyncio.run_coroutine_threadsafe( asyncio.run_coroutine_threadsafe(
self._update_download_progress_reaction(message, progress), self._update_download_progress_reaction(message, progress), loop
loop
) )
except Exception as e: except Exception as e:
logger.error(f"Error in progress callback: {e}") logger.error(f"Error in progress callback: {e}")
return progress_callback return progress_callback
async def _download_video( async def _download_video(
self, self,
downloader: DownloadManager, downloader: DownloadManager,
url: str, url: str,
progress_callback: Callable[[float], None] progress_callback: Callable[[float], None],
) -> Tuple[bool, Optional[str], Optional[str]]: ) -> Tuple[bool, Optional[str], Optional[str]]:
""" """
Download video with progress tracking. Download video with progress tracking.
@@ -422,13 +417,12 @@ class QueueHandler:
try: try:
success, file_path, error = await asyncio.wait_for( success, file_path, error = await asyncio.wait_for(
download_task, download_task, timeout=self.DOWNLOAD_TIMEOUT
timeout=self.DOWNLOAD_TIMEOUT
) )
if success: if success:
self.progress_tracker.complete_download(url) progress_tracker.complete_download(url)
else: else:
self.progress_tracker.increment_download_retries(url) progress_tracker.increment_download_retries(url)
return success, file_path, error return success, file_path, error
except asyncio.TimeoutError: except asyncio.TimeoutError:
@@ -466,7 +460,9 @@ class QueueHandler:
except asyncio.CancelledError: except asyncio.CancelledError:
pass pass
except Exception as e: except Exception as e:
logger.error(f"Error cancelling download task for {url}: {e}") logger.error(
f"Error cancelling download task for {url}: {e}"
)
self._active_downloads.clear() self._active_downloads.clear()
self._stats["active_downloads"] = 0 self._stats["active_downloads"] = 0
@@ -492,12 +488,12 @@ class QueueHandler:
logger.info("QueueHandler force cleanup completed") logger.info("QueueHandler force cleanup completed")
except Exception as e: except Exception as e:
logger.error(f"Error during QueueHandler force cleanup: {str(e)}", exc_info=True) logger.error(
f"Error during QueueHandler force cleanup: {str(e)}", exc_info=True
)
async def _update_download_progress_reaction( async def _update_download_progress_reaction(
self, self, message: discord.Message, progress: float
message: discord.Message,
progress: float
) -> None: ) -> None:
"""Update download progress reaction on message""" """Update download progress reaction on message"""
if not message: if not message:
@@ -543,9 +539,13 @@ class QueueHandler:
# Check if any downloads are stuck # Check if any downloads are stuck
current_time = datetime.utcnow() current_time = datetime.utcnow()
for url, task in self._active_downloads.items(): for url, task in self._active_downloads.items():
if not task.done() and task.get_coro().cr_frame.f_locals.get('start_time'): if not task.done() and task.get_coro().cr_frame.f_locals.get(
start_time = task.get_coro().cr_frame.f_locals['start_time'] "start_time"
if (current_time - start_time).total_seconds() > self.DOWNLOAD_TIMEOUT: ):
start_time = task.get_coro().cr_frame.f_locals["start_time"]
if (
current_time - start_time
).total_seconds() > self.DOWNLOAD_TIMEOUT:
self._stats["is_healthy"] = False self._stats["is_healthy"] = False
return False return False

View File

@@ -11,7 +11,7 @@ from pathlib import Path
from .verification_manager import VideoVerificationManager from .verification_manager import VideoVerificationManager
from .compression_manager import CompressionManager from .compression_manager import CompressionManager
from .progress_tracker import ProgressTracker from . import progress_tracker
logger = logging.getLogger("DownloadManager") logger = logging.getLogger("DownloadManager")
@@ -62,7 +62,6 @@ class DownloadManager:
# Initialize components # Initialize components
self.verification_manager = VideoVerificationManager(ffmpeg_mgr) self.verification_manager = VideoVerificationManager(ffmpeg_mgr)
self.compression_manager = CompressionManager(ffmpeg_mgr, max_file_size) self.compression_manager = CompressionManager(ffmpeg_mgr, max_file_size)
self.progress_tracker = ProgressTracker()
# Create thread pool # Create thread pool
self.download_pool = ThreadPoolExecutor( self.download_pool = ThreadPoolExecutor(
@@ -135,7 +134,7 @@ class DownloadManager:
logger.info(f"Download completed: {d['filename']}") logger.info(f"Download completed: {d['filename']}")
elif d["status"] == "downloading": elif d["status"] == "downloading":
try: try:
self.progress_tracker.update_download_progress(d) progress_tracker.update_download_progress(d)
except Exception as e: except Exception as e:
logger.debug(f"Error logging progress: {str(e)}") logger.debug(f"Error logging progress: {str(e)}")
@@ -145,7 +144,7 @@ class DownloadManager:
self.ytdl_logger.cancelled = True self.ytdl_logger.cancelled = True
self.download_pool.shutdown(wait=False, cancel_futures=True) self.download_pool.shutdown(wait=False, cancel_futures=True)
await self.compression_manager.cleanup() await self.compression_manager.cleanup()
self.progress_tracker.clear_progress() progress_tracker.clear_progress()
async def force_cleanup(self) -> None: async def force_cleanup(self) -> None:
"""Force cleanup of all resources""" """Force cleanup of all resources"""
@@ -153,7 +152,7 @@ class DownloadManager:
self.ytdl_logger.cancelled = True self.ytdl_logger.cancelled = True
self.download_pool.shutdown(wait=False, cancel_futures=True) self.download_pool.shutdown(wait=False, cancel_futures=True)
await self.compression_manager.force_cleanup() await self.compression_manager.force_cleanup()
self.progress_tracker.clear_progress() progress_tracker.clear_progress()
async def download_video( async def download_video(
self, self,
@@ -164,7 +163,7 @@ class DownloadManager:
if self._shutting_down: if self._shutting_down:
return False, "", "Downloader is shutting down" return False, "", "Downloader is shutting down"
self.progress_tracker.start_download(url) progress_tracker.start_download(url)
try: try:
# Download video # Download video
@@ -186,7 +185,7 @@ class DownloadManager:
return False, "", str(e) return False, "", str(e)
finally: finally:
self.progress_tracker.end_download(url) progress_tracker.end_download(url)
async def _safe_download( async def _safe_download(
self, self,

View File

@@ -3,11 +3,20 @@
import logging import logging
from typing import Dict, Any, Optional from typing import Dict, Any, Optional
from datetime import datetime from datetime import datetime
from enum import Enum, auto
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class ProgressStatus(Enum):
"""Status of a progress operation"""
PENDING = auto()
ACTIVE = auto()
COMPLETED = auto()
FAILED = auto()
CANCELLED = auto()
class ProgressTracker: class ProgressTracker:
"""Progress tracker singleton.""" """Progress tracker for downloads and compressions."""
_instance = None _instance = None
def __new__(cls): def __new__(cls):
@@ -18,296 +27,71 @@ class ProgressTracker:
def __init__(self): def __init__(self):
if not hasattr(self, '_initialized'): if not hasattr(self, '_initialized'):
self._data: Dict[str, Dict[str, Any]] = {} self._download_progress: Dict[str, Dict[str, Any]] = {}
self._compression_progress: Dict[str, Dict[str, Any]] = {}
self._initialized = True self._initialized = True
def update(self, key: str, data: Dict[str, Any]) -> None: def update_download_progress(self, url: str, data: Dict[str, Any]) -> None:
"""Update progress for a key.""" """Update progress for a download."""
if key not in self._data: if url not in self._download_progress:
self._data[key] = { self._download_progress[url] = {
'active': True,
'start_time': datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
'percent': 0,
'retries': 0
}
self._download_progress[url].update(data)
self._download_progress[url]['last_update'] = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
logger.debug(f"Download progress for {url}: {self._download_progress[url].get('percent', 0)}%")
def increment_download_retries(self, url: str) -> None:
"""Increment retry count for a download."""
if url in self._download_progress:
self._download_progress[url]['retries'] = self._download_progress[url].get('retries', 0) + 1
logger.debug(f"Incremented retries for {url} to {self._download_progress[url]['retries']}")
def get_download_progress(self, url: Optional[str] = None) -> Dict[str, Any]:
"""Get progress for a download."""
if url is None:
return self._download_progress
return self._download_progress.get(url, {})
def update_compression_progress(self, file_path: str, data: Dict[str, Any]) -> None:
"""Update progress for a compression."""
if file_path not in self._compression_progress:
self._compression_progress[file_path] = {
'active': True, 'active': True,
'start_time': datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"), 'start_time': datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
'percent': 0 'percent': 0
} }
self._data[key].update(data) self._compression_progress[file_path].update(data)
self._data[key]['last_update'] = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") self._compression_progress[file_path]['last_update'] = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
logger.debug(f"Progress for {key}: {self._data[key].get('percent', 0)}%") logger.debug(f"Compression progress for {file_path}: {self._compression_progress[file_path].get('percent', 0)}%")
def get(self, key: Optional[str] = None) -> Dict[str, Any]: def get_compression_progress(self, file_path: Optional[str] = None) -> Dict[str, Any]:
"""Get progress for a key.""" """Get progress for a compression."""
if key is None:
return self._data
return self._data.get(key, {})
def complete(self, key: str) -> None:
"""Mark progress as complete."""
if key in self._data:
self._data[key]['active'] = False
logger.info(f"Operation completed for {key}")
def clear(self) -> None:
"""Clear all progress data."""
self._data.clear()
logger.info("Progress data cleared")
_tracker = ProgressTracker()
def get_compression(self, file_path: Optional[str] = None) -> Dict[str, Any]:
"""Get compression progress."""
if file_path is None: if file_path is None:
return self._compressions return self._compression_progress
return self._compressions.get(file_path, {}) return self._compression_progress.get(file_path, {})
def complete_download(self, url: str) -> None: def complete_download(self, url: str) -> None:
"""Mark download as complete.""" """Mark download as complete."""
if url in self._downloads: if url in self._download_progress:
self._downloads[url]['active'] = False self._download_progress[url]['active'] = False
logger.info(f"Download completed for {url}") logger.info(f"Download completed for {url}")
def complete_compression(self, file_path: str) -> None: def complete_compression(self, file_path: str) -> None:
"""Mark compression as complete.""" """Mark compression as complete."""
if file_path in self._compressions:
self._compressions[file_path]['active'] = False
logger.info(f"Compression completed for {file_path}")
def clear(self) -> None:
"""Clear all progress data."""
self._downloads.clear()
self._compressions.clear()
logger.info("Progress data cleared")
# Global instance
_tracker = ProgressTrack
# Global instance
_tracker = ProgressTracker()
def get_tracker() -> Progre
"""Clear all progress tracking"""
self._download_progress.clear()
self._compression_progress.clear()
logger.info("Cleared all progress tracking data")
# Create singleton instance
progress_tracker = ProgressTracker()
def get_progress_tracker() -> ProgressTracker:
def mark_compression_complete(self, file_path: str) -> None:
"""Mark a compression operation as complete"""
if file_path in self._compression_progress: if file_path in self._compression_progress:
self._compression_progress[file_path]['active'] = False self._compression_progress[file_path]['active'] = False
logger.info(f"Compression completed for {file_path}") logger.info(f"Compression completed for {file_path}")
def clear_progress(self) -> None: def clear(self) -> None:
"""Clear all progress tracking""" """Clear all progress data."""
self._download_progress.clear() self._download_progress.clear()
self._compression_progress.clear() self._compression_progress.clear()
logger.info("Cleared all progress tracking data") logger.info("Progress data cleared")
# Create singleton instance def is_healthy(self) -> bool:
progress_tracker = ProgressTracker() """Check if tracker is healthy."""
return True # Basic health check, can be expanded if needed
# Export the singleton instance
def get_progress_tracker() -> ProgressTracker:
Args:
data: Dictionary containing download progress data
"""
try:
info_dict = data.get("info_dict", {})
url = info_dict.get("webpage_url")
if not url or url not in self._download_progress:
return
if data.get("status") == "downloading":
percent_str = data.get("_percent_str", "0").replace("%", "")
try:
percent = float(percent_str)
except ValueError:
percent = 0.0
total_bytes = (
data.get("total_bytes", 0) or
data.get("total_bytes_estimate", 0)
)
self._download_progress[url].update({
"active": True,
"percent": percent,
"speed": data.get("_speed_str", "N/A"),
"eta": data.get("_eta_str", "N/A"),
"downloaded_bytes": data.get("downloaded_bytes", 0),
"total_bytes": total_bytes,
"retries": data.get("retry_count", 0),
"fragment_count": data.get("fragment_count", 0),
"fragment_index": data.get("fragment_index", 0),
"video_title": info_dict.get("title", "Unknown"),
"extractor": info_dict.get("extractor", "Unknown"),
"format": info_dict.get("format", "Unknown"),
"resolution": info_dict.get("resolution", "Unknown"),
"fps": info_dict.get("fps", "Unknown"),
"last_update": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
})
logger.debug(
f"Download progress for {url}: "
f"{percent:.1f}% at {self._download_progress[url]['speed']}, "
f"ETA: {self._download_progress[url]['eta']}"
)
except Exception as e:
logger.error(f"Error updating download progress: {e}", exc_info=True)
def end_download(self, url: str, status: ProgressStatus = ProgressStatus.COMPLETED) -> None:
"""
Mark a download as completed.
Args:
url: The URL being downloaded
status: The final status of the download
"""
if url in self._download_progress:
self._download_progress[url]["active"] = False
logger.info(f"Download {status.value} for {url}")
def start_compression(self, params: CompressionParams) -> None:
"""
Initialize progress tracking for compression.
Args:
params: Compression parameters
"""
current_time = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
self._compression_progress[params.input_file] = CompressionProgress(
active=True,
filename=params.input_file,
start_time=current_time,
percent=0.0,
elapsed_time="0:00",
input_size=params.input_size,
current_size=0,
target_size=params.target_size,
codec=params.codec_params.get("c:v", "unknown"),
hardware_accel=params.use_hardware,
preset=params.codec_params.get("preset", "unknown"),
crf=params.codec_params.get("crf", "unknown"),
duration=params.duration,
bitrate=params.codec_params.get("b:v", "unknown"),
audio_codec=params.codec_params.get("c:a", "unknown"),
audio_bitrate=params.codec_params.get("b:a", "unknown"),
last_update=current_time,
current_time=None
)
def update_compression_progress(
self,
input_file: str,
progress: float,
elapsed_time: str,
current_size: int,
current_time: float
) -> None:
"""
Update compression progress information.
Args:
input_file: The input file being compressed
progress: Current progress percentage (0-100)
elapsed_time: Time elapsed as string
current_size: Current file size in bytes
current_time: Current timestamp in seconds
"""
if input_file in self._compression_progress:
self._compression_progress[input_file].update({
"percent": max(0.0, min(100.0, progress)),
"elapsed_time": elapsed_time,
"current_size": current_size,
"current_time": current_time,
"last_update": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
})
logger.debug(
f"Compression progress for {input_file}: "
f"{progress:.1f}%, Size: {current_size}/{self._compression_progress[input_file]['target_size']} bytes"
)
def end_compression(
self,
input_file: str,
status: ProgressStatus = ProgressStatus.COMPLETED
) -> None:
"""
Mark a compression operation as completed.
Args:
input_file: The input file being compressed
status: The final status of the compression
"""
if input_file in self._compression_progress:
self._compression_progress[input_file]["active"] = False
logger.info(f"Compression {status.value} for {input_file}")
def get_download_progress(self, url: Optional[str] = None) -> Optional[DownloadProgress]:
"""
Get progress information for a download.
Args:
url: Optional URL to get progress for. If None, returns all progress.
Returns:
Progress information for the specified download or None if not found
"""
if url is None:
return self._download_progress
return self._download_progress.get(url)
def get_compression_progress(
self,
input_file: Optional[str] = None
) -> Optional[CompressionProgress]:
"""
Get progress information for a compression operation.
Args:
input_file: Optional file to get progress for. If None, returns all progress.
Returns:
Progress information for the specified compression or None if not found
"""
if input_file is None:
return self._compression_progress
return self._compression_progress.get(input_file)
def get_active_downloads(self) -> Dict[str, DownloadProgress]:
"""
Get all active downloads.
Returns:
Dictionary of active downloads and their progress
"""
return {
url: progress
for url, progress in self._download_progress.items()
if progress.get("active", False)
}
def get_active_compressions(self) -> Dict[str, CompressionProgress]:
"""
Get all active compression operations.
Returns:
Dictionary of active compressions and their progress
"""
return {
input_file: progress
for input_file, progress in self._compression_progress.items()
if progress.get("active", False)
}
def clear_progress(self) -> None:
"""Clear all progress tracking"""
self._download_progress.clear()
self._compression_progress.clear()
logger.info("Cleared