diff --git a/docs/multi.md b/docs/multi.md index 091663d..bac7e7d 100644 --- a/docs/multi.md +++ b/docs/multi.md @@ -20,7 +20,7 @@ Use `-` to pass in YAML from standard input: ```bash echo "- url: http://www.example.com" | shot-scraper multi - ``` -If you run the tool with the `-n` or `--no-clobber` option any shots where the output file aleady exists will be skipped. +If you run the tool with the `-n` or `--no-clobber` option any shots where the output file already exists will be skipped. If every shot is skipped, `shot-scraper` exits without starting any `server:` processes. You can specify a subset of screenshots to take by specifying output files that you would like to create. For example, to take just the shots of `one.png` and `three.png` that are defined in `shots.yml` run this: ```bash diff --git a/shot_scraper/cli.py b/shot_scraper/cli.py index 9e5a15c..88f3869 100644 --- a/shot_scraper/cli.py +++ b/shot_scraper/cli.py @@ -1,5 +1,6 @@ import base64 import secrets +import signal import socket import subprocess import sys @@ -878,6 +879,14 @@ def multi( shots = [] if not isinstance(shots, list): raise click.ClickException("YAML file must contain a list") + # Nothing to do, for example because --no-clobber skipped every shot, + # so there is no need to start any servers or launch the browser + if not any( + shot.get("sh") or shot.get("python") or "url" in shot + for shot in shots + if not _skipped_by_options(shot, noclobber, outputs) + ): + return with sync_playwright() as p: context, browser_obj = _browser_context( p, @@ -894,13 +903,7 @@ def multi( ) try: for shot in shots: - if ( - noclobber - and shot.get("output") - and pathlib.Path(shot["output"]).exists() - ): - continue - if outputs and shot.get("output") and shot.get("output") not in outputs: + if _skipped_by_options(shot, noclobber, outputs): continue # Run "sh" key if shot.get("sh"): @@ -1742,11 +1745,23 @@ def _run_python_code(code): ) from ex +def _skipped_by_options(shot, noclobber, outputs): + output = shot.get("output") + if not output: + return False + if noclobber and pathlib.Path(output).exists(): + return True + return bool(outputs) and output not in outputs + + def _start_server(server): + # Run the server in its own process group so _stop_server() can stop + # any child processes too, such as the command started by "sh -c" + popen_kwargs = {"start_new_session": True} if os.name == "posix" else {} if isinstance(server, str): - proc = subprocess.Popen(server, shell=True) + proc = subprocess.Popen(server, shell=True, **popen_kwargs) elif isinstance(server, list): - proc = subprocess.Popen(map(str, server)) + proc = subprocess.Popen(map(str, server), **popen_kwargs) else: raise click.ClickException("server: must be a string or list") return proc, server @@ -1795,7 +1810,17 @@ def _cleanup_servers(server_processes, leave_server): ) else: for process, _ in server_processes: - process.kill() + _stop_server(process) + + +def _stop_server(process): + if os.name == "posix": + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + else: + process.kill() def _record_storyboard( diff --git a/tests/test_shot_scraper.py b/tests/test_shot_scraper.py index 7d80404..6c639b0 100644 --- a/tests/test_shot_scraper.py +++ b/tests/test_shot_scraper.py @@ -61,6 +61,71 @@ def test_multi_server(yaml): assert pathlib.Path("output.png").exists() +def _port_is_open(port): + try: + with socket.create_connection(("127.0.0.1", port), timeout=1): + return True + except OSError: + return False + + +@pytest.mark.skipif(sys.platform == "win32", reason="Uses POSIX shell syntax") +def test_multi_server_string_child_process_is_stopped(): + # The trailing command stops the shell from exec-ing python, so python + # runs as a child of "sh -c" - that child should still be stopped + port = find_free_port() + runner = CliRunner() + with runner.isolated_filesystem(): + open("server.yaml", "w").write( + f'- server: "python -m http.server {port}; echo stopped"\n' + f"- url: http://127.0.0.1:{port}/\n" + f" output: output.png\n" + ) + result = runner.invoke(cli, ["multi", "server.yaml"]) + assert result.exit_code == 0, result.output + assert pathlib.Path("output.png").exists() + deadline = time.monotonic() + 5 + while _port_is_open(port) and time.monotonic() < deadline: + time.sleep(0.1) + assert not _port_is_open(port) + + +SKIPPED_SERVER_YAML = """ +- server: python -m http.server 9024 +- url: http://localhost:9024/ + output: output.png +""".strip() + + +@pytest.mark.parametrize("options", (["-n"], ["-o", "other.png"])) +def test_multi_nothing_to_do_skips_server_and_browser(mocker, options): + start_server = mocker.patch.object(cli_module, "_start_server") + sync_playwright = mocker.patch.object(cli_module, "sync_playwright") + runner = CliRunner() + with runner.isolated_filesystem(): + open("output.png", "w").write("existing") + open("shots.yml", "w").write(SKIPPED_SERVER_YAML) + result = runner.invoke(cli, ["multi", "shots.yml", *options]) + assert result.exit_code == 0, result.output + assert open("output.png").read() == "existing" + start_server.assert_not_called() + sync_playwright.assert_not_called() + + +def test_multi_no_clobber_starts_server_for_missing_shot(): + runner = CliRunner() + with runner.isolated_filesystem(): + open("existing.png", "w").write("existing") + open("shots.yml", "w").write( + SKIPPED_SERVER_YAML + + "\n- url: http://localhost:9024/\n output: existing.png\n" + ) + result = runner.invoke(cli, ["multi", "shots.yml", "-n"]) + assert result.exit_code == 0, result.output + assert pathlib.Path("output.png").exists() + assert open("existing.png").read() == "existing" + + def test_multi_server_slow_start(): # https://github.com/simonw/shot-scraper/issues/197 # A server that takes longer than one second to start listening @@ -879,10 +944,6 @@ def __enter__(self): def __exit__(self, exc_type, exc, tb): pass - class FakeServerProcess: - def kill(self): - events.append("server.kill") - fake_page = FakePage() fake_context = FakeContext(fake_page) fake_browser = FakeBrowser() @@ -906,9 +967,14 @@ def kill(self): cli_module, "_start_server", side_effect=lambda server: ( - events.append(("server", server)) or (FakeServerProcess(), server) + events.append(("server", server)) or (object(), server) ), ) + mocker.patch.object( + cli_module, + "_stop_server", + side_effect=lambda process: events.append("server.kill"), + ) mocker.patch.object( cli_module.time, "sleep", @@ -995,9 +1061,6 @@ def __exit__(self, exc_type, exc, tb): pass class FakeServerProcess: - def kill(self): - events.append("server.kill") - def poll(self): return None @@ -1015,6 +1078,11 @@ def poll(self): events.append(("server", server)) or (FakeServerProcess(), server) ), ) + mocker.patch.object( + cli_module, + "_stop_server", + side_effect=lambda process: events.append("server.kill"), + ) mocker.patch.object( cli_module.time, "sleep",