From fae9babee52fc109c643501dd74cb9f75d18d19b Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 24 Mar 2024 21:35:16 -0700
Subject: [PATCH] --init-script prototype, refs #147

---
 shot_scraper/cli.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/shot_scraper/cli.py b/shot_scraper/cli.py
index 3f1245e..85b60df 100644
--- a/shot_scraper/cli.py
+++ b/shot_scraper/cli.py
@@ -245,6 +245,12 @@ def cli():
 @bypass_csp_option
 @silent_option
 @http_auth_options
+@click.option(
+    "init_scripts",
+    "--init-script",
+    multiple=True,
+    help="JavaScript to run as a page init script",
+)
 def shot(
     url,
     auth,
@@ -278,6 +284,7 @@ def shot(
     silent,
     auth_username,
     auth_password,
+    init_scripts,
 ):
     """
     Take a single screenshot of a page or portion of a page.
@@ -367,6 +374,7 @@ def shot(
                     log_requests=log_requests,
                     log_console=log_console,
                     silent=silent,
+                    init_scripts=init_scripts,
                 )
                 sys.stdout.buffer.write(shot)
             else:
@@ -380,6 +388,7 @@ def shot(
                     skip=skip,
                     fail=fail,
                     silent=silent,
+                    init_scripts=init_scripts,
                 )
         except TimeoutError as e:
             raise click.ClickException(str(e))
@@ -661,6 +670,12 @@ def accessibility(
 @skip_fail_options
 @bypass_csp_option
 @http_auth_options
+@click.option(
+    "init_scripts",
+    "--init-script",
+    multiple=True,
+    help="JavaScript to run as a page init script",
+)
 def javascript(
     url,
     javascript,
@@ -678,6 +693,7 @@ def javascript(
     bypass_csp,
     auth_username,
     auth_password,
+    init_scripts,
 ):
     """
     Execute JavaScript against the page and return the result as JSON
@@ -720,6 +736,8 @@ def javascript(
             auth_password=auth_password,
         )
         page = context.new_page()
+        for init_script in init_scripts:
+            page.add_init_script(script=init_script)
         if log_console:
             page.on("console", console_log)
         response = page.goto(url)
@@ -1053,6 +1071,7 @@ def _check_and_absolutize(filepath):
         # On Windows, instantiating a Path object on `http://` or `https://` will raise an exception
         return False
 
+
 def _get_viewport(width, height):
     if width or height:
         return {
@@ -1062,6 +1081,7 @@ def _get_viewport(width, height):
     else:
         return {}
 
+
 def take_shot(
     context_or_page,
     shot,
@@ -1072,6 +1092,7 @@ def take_shot(
     skip=False,
     fail=False,
     silent=False,
+    init_scripts=None,
 ):
     url = shot.get("url") or ""
     if not url:
@@ -1082,6 +1103,7 @@ def take_shot(
 
     url = url_or_file_path(url, file_exists=_check_and_absolutize)
 
+    init_scripts = init_scripts or []
     output = shot.get("output", "").strip()
     if not output and not return_bytes:
         output = filename_for_url(url, ext="png", file_exists=os.path.exists)
@@ -1141,6 +1163,9 @@ def on_response(response):
 
     full_page = not shot.get("height")
 
+    for init_script in init_scripts:
+        page.add_init_script(script=init_script)
+
     if not use_existing_page:
         # Load page and check for errors
         response = page.goto(url)