fix(cli): report unknown dataset names without a traceback

rai_toolkit/cli.py

@@ -209,9 +209,27 @@ def _cmd_datasets_list(args: argparse.Namespace) -> int:
     return 0
 
 
+def _unknown_dataset_error(slug: str) -> str | None:
+    """The registry's message for an unregistered slug, or None if it is registered.
+
+    Only the lookup is guarded, so a KeyError raised later by a registered
+    dataset's loader still surfaces as itself.
+    """
+    from rai_toolkit.examples import ExampleRegistry
+
+    try:
+        ExampleRegistry.get(slug)
+    except KeyError as e:
+        return str(e.args[0]) if e.args else f"Example dataset '{slug}' not registered."
+    return None
+
+
 def _cmd_datasets_show(args: argparse.Namespace) -> int:
     from rai_toolkit.examples import ExampleRegistry
 
+    if error := _unknown_dataset_error(args.slug):
+        print(f"error: {error}", file=sys.stderr)
+        return 2
     desc = ExampleRegistry.get(args.slug)
     info = {
         "slug": desc.slug,
@@ -249,6 +267,9 @@ def _cmd_datasets_infer_policy_expectations(args: argparse.Namespace) -> int:
     """
     from rai_toolkit.examples import ExampleRegistry
 
+    if error := _unknown_dataset_error(args.slug):
+        print(f"error: {error}", file=sys.stderr)
+        return 2
     rows = ExampleRegistry.load(args.slug, limit=args.limit)
     drafted = []
     for row in rows:

tests/test_cli.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-PackageName: rai-toolkit
 
+import json
 from pathlib import Path
 
 import pytest
@@ -87,3 +88,60 @@ def test_datasets_parser_error_exits_with_code_two(capsys) -> None:
 
     assert error.value.code == 2
     assert "invalid choice" in capsys.readouterr().err
+
+
+@pytest.mark.parametrize(
+    "command",
+    [
+        ["datasets", "show", "bbqq"],
+        ["datasets", "infer-policy-expectations", "bbqq", "--limit", "0"],
+    ],
+)
+def test_unknown_dataset_name_is_an_input_error(command: list[str], capsys) -> None:
+    exit_code = main(command)
+
+    captured = capsys.readouterr()
+    assert exit_code == 2
+    assert captured.out == ""
+    assert "error: Example dataset 'bbqq' not registered." in captured.err
+    assert "Did you mean: bbq?" in captured.err
+    assert "Traceback" not in captured.err
+
+
+def test_unknown_dataset_name_without_a_suggestion(capsys) -> None:
+    exit_code = main(["datasets", "show", "zzzzzzzzzz"])
+
+    captured = capsys.readouterr()
+    assert exit_code == 2
+    assert captured.out == ""
+    assert "Example dataset 'zzzzzzzzzz' not registered." in captured.err
+    assert "Did you mean" not in captured.err
+
+
+def test_datasets_show_reports_a_registered_dataset(capsys) -> None:
+    exit_code = main(["datasets", "show", "bbq"])
+
+    captured = capsys.readouterr()
+    assert exit_code == 0
+    assert json.loads(captured.out)["slug"] == "bbq"
+    assert captured.err == ""
+
+
+def test_infer_policy_expectations_with_limit_zero_prints_an_empty_list(capsys) -> None:
+    exit_code = main(["datasets", "infer-policy-expectations", "bbq", "--limit", "0"])
+
+    captured = capsys.readouterr()
+    assert exit_code == 0
+    assert json.loads(captured.out) == []
+
+
+def test_loader_key_error_for_a_registered_dataset_is_not_an_unknown_name(monkeypatch) -> None:
+    from rai_toolkit.examples import ExampleRegistry
+
+    def broken_load(slug: str, limit: int | None = None) -> list[dict]:
+        raise KeyError("question")
+
+    monkeypatch.setattr(ExampleRegistry, "load", staticmethod(broken_load))
+
+    with pytest.raises(KeyError, match="question"):
+        main(["datasets", "infer-policy-expectations", "bbq", "--limit", "1"])