arborist/tests/test_pi_star_code.py

"""code-py-ast@v1 π* tests.

Python source → deterministic AST S-expression. KATs mirror the
docstring's equivalence-class examples: whitespace / comment /
quote-style / numeric-formatting invariance vs identifier and
operator distinctness. Plus invalid-input cone (non-Python,
non-bytes) and determinism.

Note: this projection is projective, not invertible — the canonical
S-expression is NOT valid Python, so re-canonicalizing the output
raises (covered by an explicit test).

Backfills the test-coverage gap for arborist/pi_star/code.py
identified during the 2026-05-10 zero-coverage sweep.
"""

from __future__ import annotations

import pytest

from arborist.pi_star import get
from arborist.pi_star.protocol import PiStarError


@pytest.fixture
def ps():
    return get("code-py-ast@v1")


def _canon(ps, src: str) -> bytes:
    return ps.canonicalize(src.encode("utf-8"))


# --- registry presence + metadata ------------------------------------


def test_registry_contains_code_py_ast():
    ps = get("code-py-ast@v1")
    assert ps.name == "code-py-ast"
    assert ps.version == "v1"
    assert ps.domain == "code"


# --- positive shape --------------------------------------------------


def test_empty_module(ps):
    out = _canon(ps, "")
    # Module(body=[] type_ignores=[]) — minimal AST shape.
    assert out.startswith(b"(Module ")
    assert b"body=[]" in out


def test_single_expression(ps):
    out = _canon(ps, "1 + 2")
    assert b"BinOp" in out
    assert b"Add" in out
    assert b"value=1" in out
    assert b"value=2" in out


def test_assignment(ps):
    out = _canon(ps, "x = 1")
    assert b"Assign" in out
    assert b"'x'" in out
    assert b"value=1" in out


def test_function_def(ps):
    out = _canon(ps, "def foo(a, b):\n    return a + b\n")
    assert b"FunctionDef" in out
    assert b"'foo'" in out
    assert b"Return" in out


# --- equivalence classes preserved -----------------------------------


@pytest.mark.parametrize(
    "a,b",
    [
        # Whitespace / indentation
        ("x=1", "x = 1"),
        ("x = 1", "x  =  1"),
        # Comments
        ("x = 1", "x = 1  # comment"),
        ("# leading\nx = 1", "x = 1"),
        # Quote style
        ("x = 'foo'", 'x = "foo"'),
        # Trailing semicolons
        ("x = 1", "x = 1;"),
        # Numeric formatting (underscores in int literals)
        ("x = 1000", "x = 1_000"),
    ],
)
def test_equivalent_sources_collapse(ps, a, b):
    assert _canon(ps, a) == _canon(ps, b)


# --- equivalence classes kept distinct -------------------------------


@pytest.mark.parametrize(
    "a,b",
    [
        # Identifier names
        ("x = 1", "y = 1"),
        # Operator types
        ("a + b", "a - b"),
        ("a * b", "a / b"),
        # Argument order in calls
        ("f(a, b)", "f(b, a)"),
        # Literal vs variable
        ("x = 1", "x = 2"),
        # String vs bytes
        ("x = 'a'", "x = b'a'"),
    ],
)
def test_distinct_sources_stay_distinct(ps, a, b):
    assert _canon(ps, a) != _canon(ps, b)


# --- determinism + ordering ------------------------------------------


def test_repeated_canonicalize_is_stable(ps):
    src = "class X:\n    def f(self): return 42\n    y = [1, 2, 3]\n"
    a = _canon(ps, src)
    b = _canon(ps, src)
    c = _canon(ps, src)
    assert a == b == c


def test_field_order_is_lexical(ps):
    """_canonical_sexp sorts node._fields for cross-version stability."""
    out = _canon(ps, "x = 1")
    # In an Assign node, lexical sort puts 'targets' before 'value'.
    body_idx = out.index(b"Assign")
    targets_idx = out.index(b"targets=", body_idx)
    value_idx = out.index(b"value=", body_idx)
    assert targets_idx < value_idx


# --- source-position fields excluded ---------------------------------


def test_line_position_does_not_change_canonical(ps):
    """lineno / col_offset are not in `_fields` so they're excluded;
    the same statement at different line numbers canonicalizes the same."""
    a = _canon(ps, "x = 1")
    b = _canon(ps, "\n\nx = 1")
    assert a == b


def test_position_fields_excluded(ps):
    """Source-position metadata MUST NOT appear in canonical output."""
    out = _canon(ps, "x = 1\ny = 2\n")
    # Confirm none of the excluded fields leaked into the encoding.
    for excluded in (b"lineno=", b"col_offset=", b"end_lineno=", b"end_col_offset="):
        assert excluded not in out, f"position field {excluded!r} leaked"


# --- invalid-input cone ----------------------------------------------


def test_non_python_raises(ps):
    with pytest.raises(PiStarError, match="not valid Python"):
        ps.canonicalize(b"def (")


def test_unbalanced_paren_raises(ps):
    with pytest.raises(PiStarError, match="not valid Python"):
        ps.canonicalize(b"x = (1 + 2")


def test_non_bytes_raises(ps):
    with pytest.raises(PiStarError, match="expects bytes"):
        ps.canonicalize("x = 1")  # type: ignore[arg-type]


def test_bytearray_accepted(ps):
    """bytearray is an accepted subtype per the bytes/bytearray guard."""
    out = ps.canonicalize(bytearray(b"x = 1"))
    assert b"Assign" in out
    assert b"value=1" in out


# --- projective (not invertible) -------------------------------------


def test_canonical_output_is_not_python_source(ps):
    """Re-canonicalizing the canonical S-expression raises — by design."""
    canonical = _canon(ps, "x = 1")
    # The S-expression starts with `(Module ...)` which isn't a Python
    # expression. ast.parse rejects it.
    with pytest.raises(PiStarError, match="not valid Python"):
        ps.canonicalize(canonical)


# --- value types ------------------------------------------------------


def test_bool_serialized_as_bool_not_int(ps):
    """bool subclasses int — handler must dispatch on bool first."""
    out = _canon(ps, "x = True")
    # If bool fell through to int, we'd see `value=1`. Must see True.
    assert b"value=True" in out
    assert b"value=1" not in out


def test_none_handled(ps):
    out = _canon(ps, "x = None")
    assert b"value=None" in out


def test_float_literal(ps):
    out = _canon(ps, "x = 3.14")
    assert b"value=3.14" in out


def test_complex_literal(ps):
    out = _canon(ps, "x = 1j")
    assert b"value=1j" in out


def test_bytes_literal(ps):
    out = _canon(ps, "x = b'hi'")
    assert b"value=b'hi'" in out