diff --git a/README.md b/README.md index 4e9fc8f..acfb6f3 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # data-transform-cli -Python CLI for **YAML ↔ JSON ↔ CSV** transforms. Built for shell pipelines and local scripts. +Python CLI for **YAML ↔ JSON ↔ CSV ↔ NDJSON** transforms. Built for shell pipelines and local scripts. Streaming-friendly where practical (stdin/stdout; full documents buffered for parse safety). Requires **Python 3.11+**. @@ -22,17 +22,21 @@ After install, the `data-transform` console script is on your `PATH`. | `to-json` | Convert input to JSON | | `to-yaml` | Convert input to YAML | | `to-csv` | Convert input to CSV (array of objects or rows) | +| `to-ndjson` | Convert input to NDJSON (one JSON value per line) | | `pretty` | Pretty-print as JSON | | `minify` | Emit compact JSON | +| `select` | Extract dotted paths from parsed data | +| `pick` | Keep only listed top-level keys | +| `omit` | Drop listed top-level keys | ### Options (most commands) - `[file]` — input path; omit or use `-` for stdin - `-o, --output ` — write to a file instead of stdout -- `-f, --from ` — force input format: `json` | `yaml` | `csv` -- `to-json --minify` — compact JSON output +- `-f, --from ` — force input format: `json` | `yaml` | `csv` | `ndjson` +- `--minify` — compact JSON output (`to-json`, `select`, `pick`, `omit`) -Format is taken from `--from`, else the file extension, else simple content heuristics. +Format is taken from `--from`, else the file extension (`.ndjson` / `.jsonl` → ndjson), else simple content heuristics. ## Usage examples @@ -46,6 +50,12 @@ cat payload.json | data-transform to-yaml # JSON array to CSV data-transform to-csv records.json -o records.csv +# JSON array to NDJSON +data-transform to-ndjson records.json + +# NDJSON to JSON array +data-transform to-json events.ndjson + # CSV to pretty JSON data-transform pretty table.csv -f csv @@ -56,16 +66,48 @@ data-transform minify bulky.json printf "a: 1\n" | data-transform to-json -f yaml ``` +### Select, pick, omit + +```bash +# Single dotted path (list index supported) +echo '{"user":{"name":"Ada"},"items":[{"id":1}]}' \ + | data-transform select user.name -f json +# → "Ada" + +echo '{"items":[{"id":1},{"id":2}]}' \ + | data-transform select items.0.id -f json +# → 1 + +# Multiple paths → object map (missing → null) +echo '{"a":1}' | data-transform select a b -f json --minify +# → {"a":1,"b":null} + +# Keep / drop top-level keys (maps over arrays of objects) +echo '[{"id":1,"name":"a","note":"x"},{"id":2,"name":"b","note":"y"}]' \ + | data-transform pick id name -f json --minify +# → [{"id":1,"name":"a"},{"id":2,"name":"b"}] + +echo '{"id":1,"secret":"x","name":"a"}' \ + | data-transform omit secret -f json --minify +# → {"id":1,"name":"a"} +``` + ### Library use ```python -from data_transform import convert, parse, serialize +from data_transform import convert, parse, serialize, get_path, pick_keys, omit_keys yaml_text = convert('{"a":1}', "json", "yaml") data = parse("id,name\n1,x\n", "csv") json_text = serialize(data, "json", pretty=True) +name = get_path({"user": {"name": "Ada"}}, "user.name") +slim = pick_keys({"id": 1, "name": "a", "note": "x"}, ["id", "name"]) ``` +## NDJSON notes + +NDJSON (also known as JSON Lines / `.jsonl`) is one JSON value per non-empty line. Parsing always yields a list. Serializing a list emits one compact JSON value per line; a non-list root becomes a single line. + ## CSV notes CSV support is intentionally small and robust (RFC 4180-ish) via the Python standard library `csv` module: quoted fields, escaped quotes (`""`), and newlines inside quotes. Object arrays use a header row; all cell values are strings when parsed from CSV. @@ -77,7 +119,7 @@ pip install -e ".[dev]" pytest ``` -Tests use **pytest** with fixtures under `tests/fixtures/` for JSON/YAML/CSV round-trips and CLI smoke checks. +Tests use **pytest** with fixtures under `tests/fixtures/` for JSON/YAML/CSV/NDJSON round-trips and CLI smoke checks. ## License diff --git a/pyproject.toml b/pyproject.toml index 42fff3f..1ea297c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta" [project] name = "data-transform-cli" -version = "0.1.0" -description = "CLI for YAML ↔ JSON ↔ CSV transforms for pipelines and shell use" +version = "0.2.0" +description = "CLI for YAML ↔ JSON ↔ CSV ↔ NDJSON transforms for pipelines and shell use" readme = "README.md" license = { text = "MIT" } requires-python = ">=3.11" authors = [{ name = "rmkr-dev" }] -keywords = ["cli", "yaml", "json", "csv", "transform", "pipeline"] +keywords = ["cli", "yaml", "json", "csv", "ndjson", "transform", "pipeline"] dependencies = [ "click>=8.1", "PyYAML>=6.0", diff --git a/src/data_transform/__init__.py b/src/data_transform/__init__.py index 3e789e2..5075de0 100644 --- a/src/data_transform/__init__.py +++ b/src/data_transform/__init__.py @@ -1,7 +1,8 @@ -"""YAML ↔ JSON ↔ CSV transforms for pipelines and shell use.""" +"""YAML ↔ JSON ↔ CSV ↔ NDJSON transforms for pipelines and shell use.""" from .convert import convert, detect_format, infer_format, parse, serialize from .csv_io import escape_csv_field, parse_csv, parse_csv_rows, stringify_csv +from .shape import get_path, omit_keys, pick_keys __all__ = [ "convert", @@ -13,6 +14,9 @@ "parse_csv_rows", "stringify_csv", "escape_csv_field", + "get_path", + "pick_keys", + "omit_keys", ] -__version__ = "0.1.0" +__version__ = "0.2.0" diff --git a/src/data_transform/cli.py b/src/data_transform/cli.py index 1af1716..4edeb3f 100644 --- a/src/data_transform/cli.py +++ b/src/data_transform/cli.py @@ -4,12 +4,15 @@ import sys from pathlib import Path -from typing import Optional +from typing import Optional, Sequence import click from . import __version__ from .convert import convert, detect_format, infer_format, parse, serialize +from .shape import get_path, omit_keys, pick_keys + +FORMAT_CHOICES = ["json", "yaml", "csv", "ndjson"] def load_input(file: Optional[str]) -> tuple[str, Optional[str]]: @@ -41,6 +44,21 @@ def write_output(out: str, output_path: Optional[str]) -> None: sys.stdout.write(out) +def _split_tokens_and_file(tokens: Sequence[str]) -> tuple[list[str], Optional[str]]: + """ + Split trailing optional file from PATH/KEY tokens. + + Last token is treated as a file when it is '-' or an existing path. + """ + items = list(tokens) + if not items: + return [], None + candidate = items[-1] + if candidate == "-" or Path(candidate).is_file(): + return items[:-1], candidate + return items, None + + def _run_transform( file: Optional[str], output: Optional[str], @@ -81,7 +99,7 @@ def _run_json_view( @click.group() @click.version_option(__version__, prog_name="data-transform") def main() -> None: - """YAML ↔ JSON ↔ CSV transforms for pipelines and shell use.""" + """YAML ↔ JSON ↔ CSV ↔ NDJSON transforms for pipelines and shell use.""" @main.command("to-json") @@ -91,9 +109,9 @@ def main() -> None: "-f", "--from", "from_format", - type=click.Choice(["json", "yaml", "csv"], case_sensitive=False), + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), default=None, - help="input format: json | yaml | csv", + help="input format: json | yaml | csv | ndjson", ) @click.option("--minify", is_flag=True, default=False, help="emit compact JSON") def to_json( @@ -115,9 +133,9 @@ def to_json( "-f", "--from", "from_format", - type=click.Choice(["json", "yaml", "csv"], case_sensitive=False), + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), default=None, - help="input format: json | yaml | csv", + help="input format: json | yaml | csv | ndjson", ) def to_yaml( file: Optional[str], output: Optional[str], from_format: Optional[str] @@ -133,9 +151,9 @@ def to_yaml( "-f", "--from", "from_format", - type=click.Choice(["json", "yaml", "csv"], case_sensitive=False), + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), default=None, - help="input format: json | yaml | csv", + help="input format: json | yaml | csv | ndjson", ) def to_csv( file: Optional[str], output: Optional[str], from_format: Optional[str] @@ -144,6 +162,24 @@ def to_csv( _run_transform(file, output, from_format, "csv") +@main.command("to-ndjson") +@click.argument("file", required=False) +@click.option("-o", "--output", default=None, help="write to file instead of stdout") +@click.option( + "-f", + "--from", + "from_format", + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), + default=None, + help="input format: json | yaml | csv | ndjson", +) +def to_ndjson( + file: Optional[str], output: Optional[str], from_format: Optional[str] +) -> None: + """Convert input to NDJSON (one JSON value per line).""" + _run_transform(file, output, from_format, "ndjson") + + @main.command("pretty") @click.argument("file", required=False) @click.option("-o", "--output", default=None, help="write to file instead of stdout") @@ -151,9 +187,9 @@ def to_csv( "-f", "--from", "from_format", - type=click.Choice(["json", "yaml", "csv"], case_sensitive=False), + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), default=None, - help="input format: json | yaml | csv", + help="input format: json | yaml | csv | ndjson", ) def pretty_cmd( file: Optional[str], output: Optional[str], from_format: Optional[str] @@ -169,9 +205,9 @@ def pretty_cmd( "-f", "--from", "from_format", - type=click.Choice(["json", "yaml", "csv"], case_sensitive=False), + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), default=None, - help="input format: json | yaml | csv", + help="input format: json | yaml | csv | ndjson", ) def minify_cmd( file: Optional[str], output: Optional[str], from_format: Optional[str] @@ -180,5 +216,127 @@ def minify_cmd( _run_json_view(file, output, from_format, minify=True) +@main.command("select") +@click.argument("tokens", nargs=-1, required=True) +@click.option("-o", "--output", default=None, help="write to file instead of stdout") +@click.option( + "-f", + "--from", + "from_format", + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), + default=None, + help="input format: json | yaml | csv | ndjson", +) +@click.option("--minify", is_flag=True, default=False, help="emit compact JSON") +def select_cmd( + tokens: tuple[str, ...], + output: Optional[str], + from_format: Optional[str], + minify: bool, +) -> None: + """Extract dotted paths from parsed data as JSON. + + Usage: data-transform select PATH [PATH ...] [file] + """ + try: + path_list, input_file = _split_tokens_and_file(tokens) + if not path_list: + raise ValueError("At least one path is required") + + text, name = load_input(input_file) + src = resolve_from(text, name, from_format) + data = parse(text, src) # type: ignore[arg-type] + + if len(path_list) == 1: + result = get_path(data, path_list[0]) + else: + result = {p: get_path(data, p) for p in path_list} + + write_output( + serialize(result, "json", pretty=not minify, minify=minify), output + ) + except Exception as err: # noqa: BLE001 + click.echo(f"error: {err}", err=True) + sys.exit(1) + + +@main.command("pick") +@click.argument("tokens", nargs=-1, required=True) +@click.option("-o", "--output", default=None, help="write to file instead of stdout") +@click.option( + "-f", + "--from", + "from_format", + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), + default=None, + help="input format: json | yaml | csv | ndjson", +) +@click.option("--minify", is_flag=True, default=False, help="emit compact JSON") +def pick_cmd( + tokens: tuple[str, ...], + output: Optional[str], + from_format: Optional[str], + minify: bool, +) -> None: + """Keep only listed top-level keys from an object (or each object in an array). + + Usage: data-transform pick KEY [KEY ...] [file] + """ + try: + key_list, input_file = _split_tokens_and_file(tokens) + if not key_list: + raise ValueError("At least one key is required") + + text, name = load_input(input_file) + src = resolve_from(text, name, from_format) + data = parse(text, src) # type: ignore[arg-type] + result = pick_keys(data, key_list) + write_output( + serialize(result, "json", pretty=not minify, minify=minify), output + ) + except Exception as err: # noqa: BLE001 + click.echo(f"error: {err}", err=True) + sys.exit(1) + + +@main.command("omit") +@click.argument("tokens", nargs=-1, required=True) +@click.option("-o", "--output", default=None, help="write to file instead of stdout") +@click.option( + "-f", + "--from", + "from_format", + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), + default=None, + help="input format: json | yaml | csv | ndjson", +) +@click.option("--minify", is_flag=True, default=False, help="emit compact JSON") +def omit_cmd( + tokens: tuple[str, ...], + output: Optional[str], + from_format: Optional[str], + minify: bool, +) -> None: + """Drop listed top-level keys from an object (or each object in an array). + + Usage: data-transform omit KEY [KEY ...] [file] + """ + try: + key_list, input_file = _split_tokens_and_file(tokens) + if not key_list: + raise ValueError("At least one key is required") + + text, name = load_input(input_file) + src = resolve_from(text, name, from_format) + data = parse(text, src) # type: ignore[arg-type] + result = omit_keys(data, key_list) + write_output( + serialize(result, "json", pretty=not minify, minify=minify), output + ) + except Exception as err: # noqa: BLE001 + click.echo(f"error: {err}", err=True) + sys.exit(1) + + if __name__ == "__main__": main() diff --git a/src/data_transform/convert.py b/src/data_transform/convert.py index 4fe82e0..b22d474 100644 --- a/src/data_transform/convert.py +++ b/src/data_transform/convert.py @@ -11,7 +11,7 @@ from .csv_io import parse_csv, stringify_csv -FormatName = Literal["json", "yaml", "csv"] +FormatName = Literal["json", "yaml", "csv", "ndjson"] def detect_format(filename: str | None) -> FormatName | None: @@ -21,6 +21,8 @@ def detect_format(filename: str | None) -> FormatName | None: lower = str(filename).lower() # Support bare names and paths. name = Path(lower).name + if name.endswith(".ndjson") or name.endswith(".jsonl"): + return "ndjson" if name.endswith(".json"): return "json" if name.endswith(".yaml") or name.endswith(".yml"): @@ -38,9 +40,22 @@ def parse(text: str, format: FormatName) -> Any: return yaml.safe_load(text) if format == "csv": return parse_csv(text) + if format == "ndjson": + return _parse_ndjson(text) raise ValueError(f"Unsupported parse format: {format}") +def _parse_ndjson(text: str) -> list[Any]: + """Parse newline-delimited JSON into a list of values.""" + values: list[Any] = [] + for line in text.splitlines(): + stripped = line.strip() + if not stripped: + continue + values.append(json.loads(stripped)) + return values + + def serialize( data: Any, format: FormatName, @@ -69,9 +84,21 @@ def serialize( if not isinstance(data, list): raise ValueError("CSV output requires an array (of objects or rows)") return stringify_csv(data) + if format == "ndjson": + return _serialize_ndjson(data) raise ValueError(f"Unsupported serialize format: {format}") +def _serialize_ndjson(data: Any) -> str: + """Emit compact NDJSON: one JSON value per line.""" + if isinstance(data, list): + lines = [ + json.dumps(item, separators=(",", ":"), ensure_ascii=False) for item in data + ] + return ("\n".join(lines) + "\n") if lines else "" + return json.dumps(data, separators=(",", ":"), ensure_ascii=False) + "\n" + + def convert( text: str, from_format: FormatName, @@ -90,6 +117,9 @@ def infer_format(text: str) -> FormatName: trimmed = text.lstrip() if not trimmed: return "json" + # Prefer NDJSON when multiple non-empty lines each look like JSON values. + if _looks_like_ndjson(trimmed): + return "ndjson" if trimmed[0] in "{[": try: json.loads(text) @@ -104,3 +134,20 @@ def infer_format(text: str) -> FormatName: ): return "csv" return "yaml" + + +def _looks_like_ndjson(text: str) -> bool: + """True when two or more non-empty lines each parse as a JSON object/array.""" + lines = [ln.strip() for ln in text.splitlines() if ln.strip()] + if len(lines) < 2: + return False + for line in lines: + if not line or line[0] not in "{[": + return False + try: + val = json.loads(line) + except json.JSONDecodeError: + return False + if not isinstance(val, (dict, list)): + return False + return True diff --git a/src/data_transform/shape.py b/src/data_transform/shape.py new file mode 100644 index 0000000..78030af --- /dev/null +++ b/src/data_transform/shape.py @@ -0,0 +1,87 @@ +"""Path extraction and object key shaping helpers.""" + +from __future__ import annotations + +from typing import Any, Mapping, Sequence + + +def _parse_path(path: str) -> list[str | int]: + """Split a dotted path into segments; raise ValueError on invalid syntax.""" + if not isinstance(path, str) or path == "": + raise ValueError(f"Invalid path: {path!r}") + parts = path.split(".") + if any(p == "" for p in parts): + raise ValueError(f"Invalid path: {path!r}") + segments: list[str | int] = [] + for part in parts: + if part.isdigit(): + segments.append(int(part)) + else: + # Reject leading-minus negatives. + if part.startswith("-") and part[1:].isdigit(): + raise ValueError(f"Invalid path segment (negative index): {part!r}") + segments.append(part) + return segments + + +def get_path(data: Any, path: str) -> Any: + """ + Resolve a dotted path against data. + + Integer segments index lists. Missing keys/indices return None. + Invalid path syntax raises ValueError. + """ + segments = _parse_path(path) + current: Any = data + for seg in segments: + if isinstance(seg, int): + if not isinstance(current, Sequence) or isinstance(current, (str, bytes)): + return None + if seg < 0 or seg >= len(current): + return None + current = current[seg] + else: + if not isinstance(current, Mapping): + return None + if seg not in current: + return None + current = current[seg] + return current + + +def pick_keys(data: Any, keys: Sequence[str]) -> Any: + """ + Keep only listed top-level keys from an object, or from each object in a list. + + Non-mapping items in a list are left unchanged. A non-mapping, non-list + root is returned unchanged. + """ + key_list = list(keys) + if isinstance(data, list): + return [_pick_one(item, key_list) for item in data] + return _pick_one(data, key_list) + + +def omit_keys(data: Any, keys: Sequence[str]) -> Any: + """ + Drop listed top-level keys from an object, or from each object in a list. + + Non-mapping items in a list are left unchanged. A non-mapping, non-list + root is returned unchanged. + """ + drop = set(keys) + if isinstance(data, list): + return [_omit_one(item, drop) for item in data] + return _omit_one(data, drop) + + +def _pick_one(item: Any, keys: Sequence[str]) -> Any: + if not isinstance(item, Mapping): + return item + return {k: item[k] for k in keys if k in item} + + +def _omit_one(item: Any, drop: set[str]) -> Any: + if not isinstance(item, Mapping): + return item + return {k: v for k, v in item.items() if k not in drop} diff --git a/tests/test_cli.py b/tests/test_cli.py index 956dec2..b677fdc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -64,3 +64,74 @@ def test_to_json_minify_emits_compact_output(): assert code == 0, stderr assert "\n " not in stdout assert json.loads(stdout) + + +def test_to_ndjson_from_json_array(): + code, stdout, stderr = run_cli( + ["to-ndjson", "-f", "json"], '[{"a":1},{"b":2}]' + ) + assert code == 0, stderr + lines = [ln for ln in stdout.splitlines() if ln] + assert lines == ['{"a":1}', '{"b":2}'] + + +def test_to_json_from_ndjson_stdin(): + code, stdout, stderr = run_cli( + ["to-json", "-f", "ndjson", "--minify"], '{"a":1}\n{"b":2}\n' + ) + assert code == 0, stderr + assert json.loads(stdout) == [{"a": 1}, {"b": 2}] + + +def test_select_single_path(): + code, stdout, stderr = run_cli( + ["select", "user.name", "-f", "json", "--minify"], + '{"user":{"name":"Ada"}}', + ) + assert code == 0, stderr + assert json.loads(stdout) == "Ada" + + +def test_select_multiple_paths_and_missing(): + code, stdout, stderr = run_cli( + ["select", "a", "b", "-f", "json", "--minify"], + '{"a":1}', + ) + assert code == 0, stderr + assert json.loads(stdout) == {"a": 1, "b": None} + + +def test_select_list_index_path(): + code, stdout, stderr = run_cli( + ["select", "items.0.id", "-f", "json", "--minify"], + '{"items":[{"id":7},{"id":8}]}', + ) + assert code == 0, stderr + assert json.loads(stdout) == 7 + + +def test_select_invalid_path_exits_1(): + code, stdout, stderr = run_cli( + ["select", "a..b", "-f", "json"], + '{"a":{"b":1}}', + ) + assert code == 1 + assert "error:" in stderr + + +def test_pick_keys_from_object(): + code, stdout, stderr = run_cli( + ["pick", "id", "name", "-f", "json", "--minify"], + '{"id":1,"name":"a","note":"x"}', + ) + assert code == 0, stderr + assert json.loads(stdout) == {"id": 1, "name": "a"} + + +def test_omit_keys_from_array(): + code, stdout, stderr = run_cli( + ["omit", "note", "-f", "json", "--minify"], + '[{"id":1,"note":"x"},{"id":2,"note":"y"}]', + ) + assert code == 0, stderr + assert json.loads(stdout) == [{"id": 1}, {"id": 2}] diff --git a/tests/test_convert.py b/tests/test_convert.py index ad56702..e1cd19b 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -78,3 +78,34 @@ def test_pretty_and_minify_json(): def test_csv_serialize_rejects_non_arrays(): with pytest.raises(ValueError, match="array"): serialize({"a": 1}, "csv") + + +def test_detects_ndjson_extensions(): + assert detect_format("events.ndjson") == "ndjson" + assert detect_format("events.jsonl") == "ndjson" + assert detect_format("events.JSONL") == "ndjson" + + +def test_infers_ndjson_from_multiline_json(): + text = '{"a":1}\n{"b":2}\n' + assert infer_format(text) == "ndjson" + # Single JSON object still json + assert infer_format('{"a":1}') == "json" + + +def test_ndjson_parse_and_serialize_roundtrip(): + text = '{"a":1}\n{"b":2}\n' + data = parse(text, "ndjson") + assert data == [{"a": 1}, {"b": 2}] + out = serialize(data, "ndjson") + assert out == '{"a":1}\n{"b":2}\n' + assert parse(out, "ndjson") == data + + +def test_ndjson_serialize_non_list(): + assert serialize({"x": 1}, "ndjson") == '{"x":1}\n' + + +def test_convert_json_array_to_ndjson(): + out = convert('[{"a":1},{"b":2}]', "json", "ndjson") + assert out == '{"a":1}\n{"b":2}\n' diff --git a/tests/test_shape.py b/tests/test_shape.py new file mode 100644 index 0000000..1baecbb --- /dev/null +++ b/tests/test_shape.py @@ -0,0 +1,63 @@ +import pytest + +from data_transform.shape import get_path, omit_keys, pick_keys + + +def test_get_path_nested_object(): + data = {"user": {"name": "Ada", "age": 36}} + assert get_path(data, "user.name") == "Ada" + assert get_path(data, "user.age") == 36 + + +def test_get_path_list_index(): + data = {"items": [{"id": 1}, {"id": 2}]} + assert get_path(data, "items.0.id") == 1 + assert get_path(data, "items.1.id") == 2 + + +def test_get_path_missing_returns_none(): + data = {"a": 1} + assert get_path(data, "b") is None + assert get_path(data, "a.x") is None + assert get_path([1, 2], "5") is None + + +def test_get_path_invalid_syntax(): + with pytest.raises(ValueError, match="Invalid path"): + get_path({}, "") + with pytest.raises(ValueError, match="Invalid path"): + get_path({}, "a..b") + with pytest.raises(ValueError, match="Invalid path"): + get_path({}, ".a") + with pytest.raises(ValueError, match="negative"): + get_path([1, 2], "-1") + + +def test_pick_keys_object(): + data = {"id": 1, "name": "a", "note": "x"} + assert pick_keys(data, ["id", "name"]) == {"id": 1, "name": "a"} + assert pick_keys(data, ["missing"]) == {} + + +def test_pick_keys_array_of_objects(): + data = [{"id": 1, "name": "a"}, {"id": 2, "name": "b", "extra": True}] + assert pick_keys(data, ["id", "name"]) == [ + {"id": 1, "name": "a"}, + {"id": 2, "name": "b"}, + ] + + +def test_omit_keys_object(): + data = {"id": 1, "secret": "x", "name": "a"} + assert omit_keys(data, ["secret"]) == {"id": 1, "name": "a"} + + +def test_omit_keys_array_of_objects(): + data = [{"id": 1, "secret": "x"}, {"id": 2, "secret": "y", "ok": True}] + assert omit_keys(data, ["secret"]) == [{"id": 1}, {"id": 2, "ok": True}] + + +def test_pick_omit_passthrough_non_objects(): + assert pick_keys(42, ["a"]) == 42 + assert omit_keys("hi", ["a"]) == "hi" + assert pick_keys([1, {"a": 1, "b": 2}], ["a"]) == [1, {"a": 1}]