[FEATURE] Add better download progress (#366)

This commit is contained in:
Jesse Bannon 2022-11-26 23:33:14 -08:00 committed by GitHub
parent a77885f4de
commit 954dc0f3ee
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 41 additions and 8 deletions

View file

@ -51,6 +51,8 @@ from ytdl_sub.validators.strict_dict_validator import StrictDictValidator
from ytdl_sub.ytdl_additions.enhanced_download_archive import DownloadArchiver from ytdl_sub.ytdl_additions.enhanced_download_archive import DownloadArchiver
from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive
# pylint: disable=too-many-instance-attributes
download_logger = Logger.get(name="downloader") download_logger = Logger.get(name="downloader")
@ -150,6 +152,9 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC):
self.parents: List[EntryParent] = [] self.parents: List[EntryParent] = []
self.downloaded_entries: Dict[str, Entry] = {} self.downloaded_entries: Dict[str, Entry] = {}
self._url_total_entries: int = 0
self._url_entries_downloaded: int = 0
@property @property
def download_ytdl_options(self) -> Dict: def download_ytdl_options(self) -> Dict:
""" """
@ -440,7 +445,6 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC):
return Entry(download_entry_dict, working_directory=self.working_directory) return Entry(download_entry_dict, working_directory=self.working_directory)
def _download_entry(self, entry: Entry) -> Entry: def _download_entry(self, entry: Entry) -> Entry:
download_logger.info("Downloading entry %s", entry.title)
download_entry = self._extract_entry_info_with_retry(entry=entry) download_entry = self._extract_entry_info_with_retry(entry=entry)
upload_date_idx = self._enhanced_download_archive.mapping.get_num_entries_with_upload_date( upload_date_idx = self._enhanced_download_archive.mapping.get_num_entries_with_upload_date(
@ -468,9 +472,23 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC):
# Download entries in reverse order since they are scraped in the opposite direction. # Download entries in reverse order since they are scraped in the opposite direction.
# Helps deal with break_on_existing # Helps deal with break_on_existing
for entry in reversed(entries): for entry in reversed(entries):
self._url_entries_downloaded += 1
if self._is_downloaded(entry): if self._is_downloaded(entry):
download_logger.info(
"Already downloaded entry %d/%d: %s",
self._url_entries_downloaded,
self._url_total_entries,
entry.title,
)
continue continue
download_logger.info(
"Downloading entry %d/%d: %s",
self._url_entries_downloaded,
self._url_total_entries,
entry.title,
)
yield self._download_entry(entry) yield self._download_entry(entry)
self._mark_downloaded(entry) self._mark_downloaded(entry)
@ -528,17 +546,12 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC):
def _download( def _download(
self, self,
parents: Optional[List[EntryParent]] = None, parents: List[EntryParent],
orphans: Optional[List[Entry]] = None, orphans: List[Entry],
) -> Generator[Entry, None, None]: ) -> Generator[Entry, None, None]:
""" """
Downloads the leaf entries from EntryParent trees Downloads the leaf entries from EntryParent trees
""" """
if parents is None:
parents = []
if orphans is None:
orphans = []
# Delete info json files afterwards so other collection URLs do not use them # Delete info json files afterwards so other collection URLs do not use them
with self._separate_download_archives(clear_info_json_files=True): with self._separate_download_archives(clear_info_json_files=True):
for parent in parents: for parent in parents:
@ -557,6 +570,15 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC):
parents, orphan_entries = self._download_url_metadata(collection_url=collection_url) parents, orphan_entries = self._download_url_metadata(collection_url=collection_url)
collection_url_entries: List[Entry] = [] collection_url_entries: List[Entry] = []
# Set url entry trackers to print progress
self._url_total_entries = sum(parent.num_children() for parent in parents) + len(
orphan_entries
)
self._url_entries_downloaded = 0
download_logger.info(
"Beginning downloads for %s", self.overrides.apply_formatter(collection_url.url)
)
for entry in self._download(parents=parents, orphans=orphan_entries): for entry in self._download(parents=parents, orphans=orphan_entries):
yield entry yield entry
collection_url_entries.append(entry) collection_url_entries.append(entry)

View file

@ -67,6 +67,17 @@ class EntryParent(BaseEntry):
[child.to_type(Entry) for child in self.child_entries if self.is_entry(child)] [child.to_type(Entry) for child in self.child_entries if self.is_entry(child)]
) )
@functools.cache
def num_children(self) -> int:
"""
Returns
-------
Total number of (nested) entry children
"""
return sum(parent.num_children() for parent in self.parent_children()) + len(
self.entry_children()
)
def _playlist_variables(self, idx: int, children: List[TBaseEntry], parent_type: str) -> Dict: def _playlist_variables(self, idx: int, children: List[TBaseEntry], parent_type: str) -> Dict:
_count = self.kwargs_get(PLAYLIST_COUNT, len(children)) _count = self.kwargs_get(PLAYLIST_COUNT, len(children))
_index = children[idx].kwargs_get(PLAYLIST_INDEX, idx + 1) _index = children[idx].kwargs_get(PLAYLIST_INDEX, idx + 1)