Add support for --sleep in importImages

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])

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 = [