diff --git a/docs/use/cli.rst b/docs/use/cli.rst index f650ec1..3db7e77 100644 --- a/docs/use/cli.rst +++ b/docs/use/cli.rst @@ -43,6 +43,25 @@ The input file can be either of the following: {"url": "https://books.toscrape.com", "productNavigation": true} +.. _cli-params: + +Shared request parameters +========================= + +.. versionadded:: VERSION + +Use ``--params``/``-p`` to set :ref:`Zyte API request parameters +` for every request: + +.. code-block:: shell + + zyte-api urls.txt -p '{"httpResponseBody": true, "geolocation": "GB"}' + +For a plain-text :ref:`input file `, these parameters replace the +default :http:`request:browserHtml` parameter. For a `JSON Lines`_ input file, +parameters set on a given line take precedence. + + .. _output-file: Output file diff --git a/tests/test_main.py b/tests/test_main.py index 99f635d..6db6479 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,7 +1,10 @@ from __future__ import annotations import json +import os import subprocess +import sys +from io import StringIO from json import JSONDecodeError from pathlib import Path from tempfile import NamedTemporaryFile @@ -11,7 +14,7 @@ import pytest from zyte_api import RequestError -from zyte_api.__main__ import _get_argument_parser, run +from zyte_api.__main__ import _get_argument_parser, read_input, run if TYPE_CHECKING: from collections.abc import Iterable @@ -194,28 +197,35 @@ async def test_run_stop_on_errors_true(mockserver): def _run( - *, input_: str, mockserver: MockServer, cli_params: Iterable[str] | None = None + *, + input_: str, + mockserver: MockServer, + cli_params: Iterable[str] | None = None, + stdin: bool = False, ) -> subprocess.CompletedProcess[bytes]: cli_params = cli_params or () with NamedTemporaryFile("w") as url_list: url_list.write(input_) url_list.flush() - # Note: Using “python -m zyte_api” instead of “zyte-api” enables - # coverage tracking to work. + # Note: Using “python -m zyte_api” instead of “zyte-api”, and pointing + # PYTHONPATH at the source tree instead of letting the subprocess use + # the installed copy, enables coverage tracking to work. return subprocess.run( [ - "python", + sys.executable, "-m", "zyte_api", "--api-key", "a", "--api-url", mockserver.urljoin("/"), - url_list.name, + "-" if stdin else url_list.name, *cli_params, ], + input=input_.encode() if stdin else None, capture_output=True, check=False, + env={**os.environ, "PYTHONPATH": str(Path(__file__).parent.parent)}, ) @@ -279,6 +289,86 @@ def test_intype_jsonl_explicit(mockserver): ) +def test_stdin(mockserver): + result = _run( + input_="https://a.example", + mockserver=mockserver, + cli_params=["-p", '{"httpResponseBody": true}'], + stdin=True, + ) + assert not result.returncode + assert b'"httpResponseBody"' in result.stdout + + +def test_params_txt(mockserver): + result = _run( + input_="https://a.example", + mockserver=mockserver, + cli_params=["-p", '{"httpResponseBody": true}'], + ) + assert not result.returncode + assert b"browserHtml" not in result.stdout + assert b'"httpResponseBody"' in result.stdout + + +def test_params_jsonl(mockserver): + result = _run( + input_='{"url": "https://a.example", "browserHtml": true}', + mockserver=mockserver, + cli_params=["-p", '{"browserHtml": false, "httpResponseBody": true}'], + ) + assert not result.returncode + assert b'"browserHtml"' in result.stdout + assert b'"httpResponseBody"' in result.stdout + + +@pytest.mark.parametrize("value", ("{", "[]")) +def test_params_invalid(value, capsys): + parser = _get_argument_parser() + with pytest.raises(SystemExit): + parser.parse_args(["--params", value, "README.rst"]) + assert "--params/-p" in capsys.readouterr().err + + +def test_read_input_txt(): + assert read_input(StringIO("https://a.example\n\n"), "txt") == [ + { + "url": "https://a.example", + "browserHtml": True, + "echoData": "https://a.example", + } + ] + + +def test_read_input_txt_params(): + parser = _get_argument_parser() + args = parser.parse_args(["-p", '{"httpResponseBody": true}', "README.rst"]) + assert read_input(StringIO("https://a.example\n"), "txt", args.params) == [ + { + "url": "https://a.example", + "httpResponseBody": True, + "echoData": "https://a.example", + } + ] + + +def test_read_input_jl_params(): + input_fp = StringIO('{"url": "https://a.example", "browserHtml": true}\n\n') + params = {"browserHtml": False, "httpResponseBody": True} + assert read_input(input_fp, "jl", params) == [ + { + "url": "https://a.example", + "browserHtml": True, + "httpResponseBody": True, + "echoData": "https://a.example", + } + ] + + +def test_read_input_empty(): + assert read_input(StringIO(""), "txt") == [] + + @pytest.mark.flaky(reruns=16) def test_limit_and_shuffle(mockserver): result = _run( diff --git a/zyte_api/__main__.py b/zyte_api/__main__.py index 824d88d..337cf0e 100644 --- a/zyte_api/__main__.py +++ b/zyte_api/__main__.py @@ -112,7 +112,9 @@ def write_output(content: Any) -> None: def read_input( - input_fp: IO[str], intype: Literal["txt", "jl"] | object + input_fp: IO[str], + intype: Literal["txt", "jl"] | object, + params: dict[str, Any] | None = None, ) -> list[dict[str, Any]]: assert intype in {"txt", "jl", _UNSET} lines = input_fp.readlines() @@ -122,9 +124,14 @@ def read_input( intype = _guess_intype(input_fp.name, lines) if intype == "txt": urls = [u.strip() for u in lines if u.strip()] - records = [{"url": url, "browserHtml": True} for url in urls] + base = params if params else {"browserHtml": True} + records = [{"url": url, **base} for url in urls] else: - records = [json.loads(line.strip()) for line in lines if line.strip()] + records = [ + {**(params or {}), **json.loads(line.strip())} + for line in lines + if line.strip() + ] # Automatically replicating the url in echoData to being able to # to match URLs with content in the responses for record in records: @@ -132,6 +139,16 @@ def read_input( return records +def _parse_params(value: str) -> dict[str, Any]: + try: + params = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"invalid JSON: {e}") from e + if not isinstance(params, dict): + raise argparse.ArgumentTypeError("expected a JSON object") + return params + + def _get_argument_parser(program_name: str = "zyte-api") -> argparse.ArgumentParser: p = argparse.ArgumentParser( prog=program_name, @@ -157,6 +174,19 @@ def _get_argument_parser(program_name: str = "zyte-api") -> argparse.ArgumentPar "with 'txt' as fallback." ), ) + p.add_argument( + "--params", + "-p", + type=_parse_params, + help=( + "JSON object of Zyte API request parameters to use for every " + "request.\n" + "\n" + "For a plain-text input file, these parameters replace the " + "default browserHtml parameter. For a JSON Lines input file, " + "parameters set on a given line take precedence." + ), + ) p.add_argument("--limit", type=int, help="Maximum number of requests to send.") p.add_argument( "--output", @@ -265,11 +295,11 @@ def _main(program_name: str = "zyte-api") -> None: if args.INPUT == "-": with nullcontext(sys.stdin) as input_fp: - queries = read_input(input_fp, args.intype) + queries = read_input(input_fp, args.intype, args.params) else: try: with Path(args.INPUT).open(encoding="utf8") as input_fp: - queries = read_input(input_fp, args.intype) + queries = read_input(input_fp, args.intype, args.params) except OSError as e: p.error(f"Cannot open input file {args.INPUT!r}: {e}") if not queries: