diff --git a/src/auto_archiver/archivers/generic_archiver/generic_archiver.py b/src/auto_archiver/archivers/generic_archiver/generic_archiver.py index e339434..59cd3f8 100644 --- a/src/auto_archiver/archivers/generic_archiver/generic_archiver.py +++ b/src/auto_archiver/archivers/generic_archiver/generic_archiver.py @@ -79,7 +79,7 @@ class GenericArchiver(Archiver): except Exception as e: logger.error(f"Error downloading cover image {thumbnail_url}: {e}") - dropin = self.dropin_for_extractor(info_extractor) + dropin = self.dropin_for_name(info_extractor.ie_key()) if dropin: try: metadata = dropin.download_additional_media(video_data, info_extractor, metadata) @@ -102,7 +102,7 @@ class GenericArchiver(Archiver): '_format_sort_fields', 'chapters', 'requested_formats', 'format_note', 'audio_channels', 'asr', 'fps', 'was_live', 'is_live', 'heatmap', 'age_limit', 'stretched_ratio'] - dropin = self.dropin_for_extractor(info_extractor) + dropin = self.dropin_for_name(info_extractor.ie_key()) if dropin: try: base_keys += dropin.keys_to_clean(video_data, info_extractor) @@ -157,7 +157,7 @@ class GenericArchiver(Archiver): ie_instance = info_extractor(downloader=ydl) post_data = None - dropin = self.dropin_for_extractor(info_extractor) + dropin = self.dropin_for_name(info_extractor.ie_key()) if not dropin: # TODO: add a proper link to 'how to create your own dropin' logger.debug(f"""Could not find valid dropin for {info_extractor.IE_NAME}. @@ -207,8 +207,7 @@ class GenericArchiver(Archiver): return self.add_metadata(data, info_extractor, url, result) - def dropin_for_extractor(self, info_extractor: Type[InfoExtractor], additional_paths = []): - dropin_name = info_extractor.ie_key().lower() + def dropin_for_name(self, dropin_name: str, additional_paths = [], package=__package__) -> Type[InfoExtractor]: if dropin_name == "generic": # no need for a dropin for the generic extractor (?) @@ -241,7 +240,7 @@ class GenericArchiver(Archiver): # fallback to loading the dropins within auto-archiver try: - return _load_dropin(importlib.import_module(f".{dropin_name}", package=__package__)) + return _load_dropin(importlib.import_module(f".{dropin_name}", package=package)) except ModuleNotFoundError: pass @@ -258,7 +257,7 @@ class GenericArchiver(Archiver): ydl.params['getcomments'] = False result = False - dropin_submodule = self.dropin_for_extractor(info_extractor) + dropin_submodule = self.dropin_for_name(info_extractor.ie_key()) try: if dropin_submodule and dropin_submodule.skip_ytdlp_download(info_extractor, url): diff --git a/tests/archivers/test_generic_archiver.py b/tests/archivers/test_generic_archiver.py index d493437..6e249e8 100644 --- a/tests/archivers/test_generic_archiver.py +++ b/tests/archivers/test_generic_archiver.py @@ -1,10 +1,12 @@ -import pytest from pathlib import Path -import datetime +import datetime import os -from auto_archiver.archivers.generic_archiver import GenericArchiver +from os.path import dirname +import pytest + +from auto_archiver.archivers.generic_archiver import GenericArchiver from .test_archiver_base import TestArchiverBase class TestGenericArchiver(TestArchiverBase): @@ -23,6 +25,17 @@ class TestGenericArchiver(TestArchiverBase): 'cookies_from_browser': False, 'cookie_file': None, } + + def test_load_dropin(self): + # test loading dropins that are in the generic_archiver package + package = "auto_archiver.archivers.generic_archiver" + assert self.archiver.dropin_for_name("bluesky", package=package) + + # test loading dropings via filepath + path = os.path.join(dirname(dirname(__file__)), "data/") + assert self.archiver.dropin_for_name("dropin", additional_paths=[path]) + + @pytest.mark.parametrize("url, is_suitable", [ ("https://www.youtube.com/watch?v=5qap5aO4i9A", True), diff --git a/tests/data/dropin.py b/tests/data/dropin.py new file mode 100644 index 0000000..0049c48 --- /dev/null +++ b/tests/data/dropin.py @@ -0,0 +1,5 @@ +# this is a dummy class used to test importing a dropin in the +# generic extractor by filename/path + +class Dropin: + pass \ No newline at end of file