diff --git a/miraheze/mediawiki/mwimport.py b/miraheze/mediawiki/mwimport.py index 6dbdb3a..9ed52b0 100644 --- a/miraheze/mediawiki/mwimport.py +++ b/miraheze/mediawiki/mwimport.py @@ -6,6 +6,12 @@ import sys +# The threshold for --images-sleep's automatic calculation, where it'll decide +# whether or not to sleep for 0s or 1s. +# https://wm-bot.wmcloud.org/logs/%23miraheze-tech-ops/20250714.txt#:~:text=[05:37:19],That's%20sensible +_IMAGES_SLEEP_AUTO_THRESHOLD = 1000 + + def parse_args(input_args: list | None = None, check_paths: bool = True) -> argparse.Namespace: parser = argparse.ArgumentParser(description='A script to automate manual wiki imports') parser.add_argument( @@ -28,6 +34,10 @@ def parse_args(input_args: list | None = None, check_paths: bool = True) -> argp '--search-recursively', action='store_true', help='Whether or not to pass --search-recursively (check files in subdirectories) to importImages.php', ) + parser.add_argument( + '--images-sleep', type=int, default=-1, + help='The time to sleep between importing images for importImages.php (negative for auto-calculation)', + ) parser.add_argument('wiki', help='Database name of the wiki to import to') args = parser.parse_args(input_args) @@ -46,9 +56,28 @@ def parse_args(input_args: list | None = None, check_paths: bool = True) -> argp if args.images and not os.path.exists(args.images): raise ValueError(f'Cannot find images to import: {repr(args.images)}') + if args.images and args.images_sleep < 0: + args.images_sleep = calculate_images_sleep(args.images) if check_paths else 0 + return args +def calculate_images_sleep(images: str) -> int: + # In the interest of code simplicity, all calculations are done assuming that + # --search-recursively is passed. It is unlikely where one wants to only upload + # files from a directory but not its subdirectories anyway, and this is meant + # to be a "eh, good enough" heuristic, so an "eh, good enough" algorithm for + # edge cases seems acceptable. + total = 0 + + for _, _, files in os.walk(images): + total += len(files) + if total >= _IMAGES_SLEEP_AUTO_THRESHOLD: + return 1 + + return 0 + + def log(message: str): # pragma: no cover subprocess.run( ['/usr/local/bin/logsalmsg', message], @@ -76,7 +105,7 @@ def get_scripts(args: argparse.Namespace) -> list[list[str]]: scripts.append(script) if args.images: - script = ['importImages', f'--comment={args.images_comment}'] + script = ['importImages', f'--sleep={args.images_sleep}', f'--comment={args.images_comment}'] if args.search_recursively: script.append('--search-recursively') script.extend(['--', args.images]) diff --git a/tests/test_mwimport.py b/tests/test_mwimport.py index 7f2bf28..6036938 100644 --- a/tests/test_mwimport.py +++ b/tests/test_mwimport.py @@ -105,6 +105,54 @@ def test_parse_args_both_xml_images_exists(): assert args.images == images +def test_parse_args_images_sleep_manual(): + args = mwimport.parse_args([ + '--images=images', + '--images-comment=Importing from https://example.com', + '--images-sleep=5', + 'examplewiki', + ], False) + + assert args.images_sleep == 5 + + +def test_parse_args_images_sleep_auto_below_threshold(): + with tempfile.TemporaryDirectory() as tempdir: + # Populate five test files + for i in range(5): + with open(os.path.join(tempdir, f'{i}.txt'), 'w'): + pass + + args = mwimport.parse_args([ + f'--images={tempdir}', + '--images-comment=Importing from https://example.com', + 'examplewiki', + ]) + + assert args.images_sleep == 0 + + +def test_parse_args_images_sleep_auto_above_threshold(): + with tempfile.TemporaryDirectory() as tempdir: + # Populate 1000 test files, bucketed in 10 directories + # (to test file counts with subdirectories) + for folder in range(10): + folder = os.path.join(tempdir, str(folder)) + os.mkdir(folder) + + for file in range(100): + with open(os.path.join(folder, f'{file}.txt'), 'w'): + pass + + args = mwimport.parse_args([ + f'--images={tempdir}', + '--images-comment=Importing from https://example.com', + 'examplewiki', + ]) + + assert args.images_sleep == 1 + + def test_get_scripts_xml_images(): args = mwimport.parse_args([ '--version=0.42', @@ -116,7 +164,7 @@ def test_get_scripts_xml_images(): scripts = mwimport.get_scripts(args) expected = [ ['importDump', '--no-updates', '--', 'dump.xml'], - ['importImages', '--comment=Importing from https://example.com', '--', 'images'], + ['importImages', '--sleep=0', '--comment=Importing from https://example.com', '--', 'images'], ['rebuildall'], ['initEditCount'], ['initSiteStats', '--update'], @@ -157,7 +205,7 @@ def test_get_scripts_search_recursively(): ], False) scripts = mwimport.get_scripts(args) expected = [ - ['importImages', '--comment=Importing from https://example.com', '--search-recursively', '--', 'images'], + ['importImages', '--sleep=0', '--comment=Importing from https://example.com', '--search-recursively', '--', 'images'], ['initSiteStats', '--update'], ] expected = [