From 92c99889878df146e17b6ef09d96ec54471493c0 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Sun, 27 Nov 2022 22:38:50 -0800 Subject: [PATCH] [BACKEND] Reduce memory footprint (#375) * [BACKEND] Reduce memory footprint * entry refactor --- src/ytdl_sub/cli/main.py | 2 + src/ytdl_sub/downloaders/downloader.py | 52 +++++++++++--------------- src/ytdl_sub/entries/base_entry.py | 8 ++++ src/ytdl_sub/entries/entry_parent.py | 36 +++++++++--------- 4 files changed, 50 insertions(+), 48 deletions(-) diff --git a/src/ytdl_sub/cli/main.py b/src/ytdl_sub/cli/main.py index 1aae4069..e5fe9e38 100644 --- a/src/ytdl_sub/cli/main.py +++ b/src/ytdl_sub/cli/main.py @@ -1,6 +1,7 @@ import argparse import errno import fcntl +import gc import os import sys import tempfile @@ -59,6 +60,7 @@ def _download_subscriptions_from_yaml_files( transaction_log = subscription.download(dry_run=args.dry_run) output.append((subscription, transaction_log)) + gc.collect() # Garbage collect after each subscription download return output diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 98c39d59..eca70886 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -30,7 +30,6 @@ from ytdl_sub.downloaders.generic.validators import MultiUrlValidator from ytdl_sub.downloaders.generic.validators import UrlThumbnailListValidator from ytdl_sub.downloaders.generic.validators import UrlValidator from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder -from ytdl_sub.entries.base_entry import BaseEntry from ytdl_sub.entries.entry import Entry from ytdl_sub.entries.entry_parent import EntryParent from ytdl_sub.entries.variables.kwargs import COMMENTS @@ -56,10 +55,6 @@ from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadAr download_logger = Logger.get(name="downloader") -def _entry_key(entry: BaseEntry) -> str: - return entry.extractor + entry.uid - - class DownloaderValidator(StrictDictValidator, AddsVariablesMixin, ABC): """ Placeholder class to define downloader options @@ -101,8 +96,6 @@ class URLDownloadState: def __init__(self, entries_total: int): self.entries_total = entries_total self.entries_downloaded = 0 - - self.entries: List[Entry] = [] self.thumbnails_downloaded: Set[str] = set() @@ -158,7 +151,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): self.overrides = overrides self._download_ytdl_options_builder = download_ytdl_options self._metadata_ytdl_options_builder = metadata_ytdl_options - self.downloaded_entries: Dict[str, Entry] = {} + self._downloaded_entries: Set[str] = set() self._url_state: Optional[URLDownloadState] = None @@ -386,10 +379,10 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): # DOWNLOAD FUNCTIONS def _is_downloaded(self, entry: Entry) -> bool: - return _entry_key(entry) in self.downloaded_entries + return entry.ytdl_uid() in self._downloaded_entries def _mark_downloaded(self, entry: Entry) -> None: - self.downloaded_entries[_entry_key(entry)] = entry + self._downloaded_entries.add(entry.ytdl_uid()) @property def collection(self) -> MultiUrlValidator: @@ -586,10 +579,8 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): ) for entry in self._download(parents=parents, orphans=orphan_entries): yield entry - # Add entry to URL state - self._url_state.entries.append(entry) # Update thumbnails in case of last_entry - self._download_url_thumbnails(collection_url=collection_url) + self._download_url_thumbnails(collection_url=collection_url, entry=entry) @classmethod def _download_thumbnail( @@ -659,27 +650,26 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): else: download_logger.warning("Failed to download thumbnail id '%s'", thumbnail_id) - def _download_url_thumbnails(self, collection_url: UrlValidator): + def _download_url_thumbnails(self, collection_url: UrlValidator, entry: Entry): """ After all media entries have been downloaded, post processed, and moved to the output directory, run this function. This lets the downloader add any extra files directly to the output directory, for things like YT channel image, banner. """ - for entry in self._url_state.entries: - if entry.kwargs_contains(PLAYLIST_ENTRY): - self._download_parent_thumbnails( - thumbnail_list_info=collection_url.playlist_thumbnails, - entry=entry, - parent=EntryParent( - entry.kwargs(PLAYLIST_ENTRY), working_directory=self.working_directory - ), - ) + if entry.kwargs_contains(PLAYLIST_ENTRY): + self._download_parent_thumbnails( + thumbnail_list_info=collection_url.playlist_thumbnails, + entry=entry, + parent=EntryParent( + entry.kwargs(PLAYLIST_ENTRY), working_directory=self.working_directory + ), + ) - if entry.kwargs_contains(SOURCE_ENTRY): - self._download_parent_thumbnails( - thumbnail_list_info=collection_url.source_thumbnails, - entry=entry, - parent=EntryParent( - entry.kwargs(SOURCE_ENTRY), working_directory=self.working_directory - ), - ) + if entry.kwargs_contains(SOURCE_ENTRY): + self._download_parent_thumbnails( + thumbnail_list_info=collection_url.source_thumbnails, + entry=entry, + parent=EntryParent( + entry.kwargs(SOURCE_ENTRY), working_directory=self.working_directory + ), + ) diff --git a/src/ytdl_sub/entries/base_entry.py b/src/ytdl_sub/entries/base_entry.py index 78a13aaf..7a714961 100644 --- a/src/ytdl_sub/entries/base_entry.py +++ b/src/ytdl_sub/entries/base_entry.py @@ -358,3 +358,11 @@ class BaseEntry(BaseEntryVariables, ABC): entry_ext = entry_dict.get("ext") return entry_ext is not None + + def ytdl_uid(self) -> str: + """ + Returns + ------- + extractor + uid, making this a unique hash for any entry + """ + return self.extractor + self.uid diff --git a/src/ytdl_sub/entries/entry_parent.py b/src/ytdl_sub/entries/entry_parent.py index 482fab43..09c7d258 100644 --- a/src/ytdl_sub/entries/entry_parent.py +++ b/src/ytdl_sub/entries/entry_parent.py @@ -1,4 +1,3 @@ -import functools import math from typing import Dict from typing import List @@ -53,21 +52,22 @@ def _sort_entries(entries: List[TBaseEntry]) -> List[TBaseEntry]: class EntryParent(BaseEntry): def __init__(self, entry_dict: Dict, working_directory: str): super().__init__(entry_dict=entry_dict, working_directory=working_directory) - self.child_entries: List["EntryParent"] = [] + self._parent_children: List["EntryParent"] = [] + self._entry_children: List[Entry] = [] - @functools.cache def parent_children(self) -> List["EntryParent"]: """This parent's children that are also parents""" - return _sort_entries([child for child in self.child_entries if self.is_entry_parent(child)]) + return self._parent_children - @functools.cache def entry_children(self) -> List[Entry]: """This parent's children that are entries""" - return _sort_entries( - [child.to_type(Entry) for child in self.child_entries if self.is_entry(child)] - ) + return self._entry_children + + @property + def children(self) -> List[TBaseEntry]: + """Children, both entries and parent entries, cast as BaseEntry""" + return self._parent_children + self._entry_children - @functools.cache def num_children(self) -> int: """ Returns @@ -166,7 +166,7 @@ class EntryParent(BaseEntry): """ Populates a tree of EntryParents that belong to this instance """ - self.child_entries = [ + entries = [ EntryParent( entry_dict=entry_dict, working_directory=self.working_directory(), @@ -174,6 +174,12 @@ class EntryParent(BaseEntry): for entry_dict in entry_dicts if entry_dict in self ] + + self._parent_children = _sort_entries([ent for ent in entries if self.is_entry_parent(ent)]) + self._entry_children = _sort_entries( + [ent.to_type(Entry) for ent in entries if self.is_entry(ent)] + ) + return self def get_thumbnail_url(self, thumbnail_id: str) -> Optional[str]: @@ -204,9 +210,7 @@ class EntryParent(BaseEntry): if not playlist_id: return False - return self.uid == playlist_id or any( - child.__contains__(item) for child in self.child_entries - ) + return self.uid == playlist_id or any(child.__contains__(item) for child in self.children) @classmethod def _get_disconnected_root_parent( @@ -220,9 +224,7 @@ class EntryParent(BaseEntry): return webpage_url in url or url in webpage_url top_level_parents = [ - parent - for parent in parents - if not parent.child_entries and _url_matches(parent.webpage_url) + parent for parent in parents if not parent.children and _url_matches(parent.webpage_url) ] match len(top_level_parents): @@ -263,7 +265,7 @@ class EntryParent(BaseEntry): # If a disconnected root parent exists, connect it here if (root_parent := cls._get_disconnected_root_parent(url, parents)) is not None: parents.remove(root_parent) - root_parent.child_entries = parents + root_parent._parent_children = parents parents = [root_parent] for parent in parents: