#!/usr/bin/env python3
"""Measure whether Python YAML parsers accept anchors/aliases, merge keys
(<<:), and multi-document streams (---).

Parsers:
  - PyYAML  yaml.safe_load / safe_load_all
  - ruamel.yaml YAML(typ="safe") load / load_all

Run:
  cd scripts/benchmarks/yaml-tool-acceptance
  python3 -m venv venv && ./venv/bin/pip install pyyaml ruamel.yaml
  ./venv/bin/python measure.py    # prints JSON, writes results-python.json
"""
import io
import json
import platform
import sys
from pathlib import Path

HERE = Path(__file__).parent

ANCHOR = "base: &b\n  x: 1\nchild: *b"
MERGE = "base: &b\n  x: 1\nchild:\n  <<: *b\n  y: 2"
MULTIDOC = "a: 1\n---\nb: 2"


def attempt(fn):
    try:
        value = fn()
        return {"outcome": "accepted", "repr": json.dumps(value, default=str)[:200]}
    except Exception as exc:  # noqa: BLE001
        return {"outcome": "error", "message": " ".join(str(exc).split())[:200]}


def main():
    import yaml as pyyaml
    import ruamel.yaml as ruamel_pkg
    from ruamel.yaml import YAML

    ruamel_safe = YAML(typ="safe")

    results = [
        {
            "parser": f"PyYAML {pyyaml.__version__}",
            "results": {
                "anchor-alias": attempt(lambda: pyyaml.safe_load(ANCHOR)),
                "merge-key": attempt(lambda: pyyaml.safe_load(MERGE)),
                "multi-doc-load": attempt(lambda: pyyaml.safe_load(MULTIDOC)),
                "multi-doc-loadAll": attempt(lambda: list(pyyaml.safe_load_all(MULTIDOC))),
            },
        },
        {
            "parser": f"ruamel.yaml {ruamel_pkg.__version__}",
            "results": {
                "anchor-alias": attempt(lambda: ruamel_safe.load(io.StringIO(ANCHOR))),
                "merge-key": attempt(lambda: ruamel_safe.load(io.StringIO(MERGE))),
                "multi-doc-load": attempt(lambda: ruamel_safe.load(io.StringIO(MULTIDOC))),
                "multi-doc-loadAll": attempt(
                    lambda: list(ruamel_safe.load_all(io.StringIO(MULTIDOC)))
                ),
            },
        },
    ]

    payload = {
        "generatedBy": "measure.py",
        "python": platform.python_version(),
        "results": results,
    }
    (HERE / "results-python.json").write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n")
    json.dump(payload, sys.stdout, indent=2, ensure_ascii=False)
    print()


if __name__ == "__main__":
    main()