pythsol 0.1.0__py3-none-any.whl

pythsol/client/gen.py

@@ -0,0 +1,496 @@
+"""Render a solders client from a :class:`ProgramModel`.
+
+The text this module returns imports only the standard library and
+``solders``. It does not import the compiler, so a generated
+``*_client.py`` can ship next to a program without a PythSol install
+beyond what the caller already uses to build transactions.
+"""
+
+from __future__ import annotations
+
+from pythsol.compiler.model import (
+    ACCOUNT_DATA_OFFSET,
+    AccountModel,
+    DataParamModel,
+    EventModel,
+    FieldModel,
+    InstructionModel,
+    ProgramModel,
+    StructModel,
+    snake_case,
+)
+
+__all__ = ["render_client"]
+
+_SYSTEM_PROGRAM = "11111111111111111111111111111111"
+_TOKEN_PROGRAM = "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"
+_TOKEN_2022_PROGRAM = "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb"
+_ATA_PROGRAM = "ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"
+
+#: ``struct.pack`` format for scalar widths that fit in one call.
+_PACK_FMT: dict[str, str] = {
+    "uint8_t": "<B",
+    "uint16_t": "<H",
+    "uint32_t": "<I",
+    "uint64_t": "<Q",
+    "int8_t": "<b",
+    "int16_t": "<h",
+    "int32_t": "<i",
+    "int64_t": "<q",
+    "bool": "<B",
+}
+
+_UNPACK_FMT: dict[str, str] = dict(_PACK_FMT)
+
+
+def render_client(model: ProgramModel) -> str:
+    """Return the full ``{name}_client.py`` source for ``model``."""
+    lines: list[str] = [
+        '"""Generated by PythSol. Do not edit.',
+        "",
+        f"Program `{model.name}` from `{model.source}`.",
+        "Build transactions with solders. Decode account bytes with `decode_*`.",
+        '"""',
+        "",
+        "from __future__ import annotations",
+        "",
+        "import struct",
+        "",
+        "from solders.instruction import AccountMeta, Instruction",
+        "from solders.pubkey import Pubkey",
+        "",
+        f"PROGRAM_NAME = {model.snake_name!r}",
+        f'_SYSTEM_PROGRAM = Pubkey.from_string("{_SYSTEM_PROGRAM}")',
+        "",
+    ]
+    if any(instr.needs_token or instr.needs_ata for instr in model.instructions):
+        lines.append(f'_TOKEN_PROGRAM = Pubkey.from_string("{_TOKEN_PROGRAM}")')
+    if any(instr.needs_token2022 for instr in model.instructions):
+        lines.append(f'_TOKEN_2022_PROGRAM = Pubkey.from_string("{_TOKEN_2022_PROGRAM}")')
+    if any(instr.needs_ata for instr in model.instructions):
+        lines.append(f'_ATA_PROGRAM = Pubkey.from_string("{_ATA_PROGRAM}")')
+    if any(
+        instr.needs_token or instr.needs_token2022 or instr.needs_ata
+        for instr in model.instructions
+    ):
+        lines.append("")
+    lines.extend(_enum_constants(model))
+    if model.errors:
+        lines.append("# Custom program errors (6000 + declaration index).")
+        for err in model.errors:
+            const = f"ERR_{err.owner.upper()}_{err.name.upper()}"
+            lines.append(f"{const} = {err.code}  # {err.owner}.{err.name}")
+        lines.append("")
+    for instr in model.instructions:
+        const = f"_DISC_{instr.snake_name.upper()}"
+        lines.append(f"{const} = {_bytes_expr(instr.discriminator)}")
+    lines.append("")
+    if _has_string_arg(model):
+        lines.extend(_fixed_string_helper())
+        lines.append("")
+    for instr in model.instructions:
+        lines.extend(_instruction_fn(instr))
+        lines.append("")
+    structs = {item.name: item for item in model.structs}
+    payloads = {item.name: item.payload_c for item in model.enums if item.payload_c}
+    for struct in model.structs:
+        lines.extend(_struct_decode_fn(struct, structs, payloads))
+        lines.append("")
+    for account in model.accounts:
+        lines.extend(_decode_fn(account, structs, payloads))
+        lines.append("")
+    if model.events:
+        lines.extend(_event_decoder(model.events))
+        lines.append("")
+    text = "\n".join(lines).rstrip() + "\n"
+    return text
+
+
+def _enum_constants(model: ProgramModel) -> list[str]:
+    """``STATUS_OPEN = 0`` for each unit-enum variant (the one-byte tag)."""
+    if not model.enums:
+        return []
+    used: set[str] = set()
+    lines = ["# Unit enum tags (declaration index, one byte)."]
+    for enum in model.enums:
+        for variant, tag in enum.variants:
+            base = f"{enum.name.upper()}_{variant.upper()}"
+            const = base if base not in used else f"{base}_{tag}"
+            used.add(const)
+            lines.append(f"{const} = {tag}  # {enum.name}.{variant}")
+    lines.append("")
+    return lines
+
+
+def _bytes_expr(data: bytes) -> str:
+    inner = ", ".join(str(byte) for byte in data)
+    return f"bytes(({inner}))"
+
+
+def _py_type(c_type: str) -> str:
+    if c_type == "SolPubkey":
+        return "Pubkey"
+    if c_type == "bool":
+        return "bool"
+    return "int"
+
+
+def _has_string_arg(model: ProgramModel) -> bool:
+    return any(
+        param.shape == "string" for instr in model.instructions for param in instr.data_params
+    )
+
+
+def _fixed_string_helper() -> list[str]:
+    """One packer for every ``FixedString[N]`` instruction argument."""
+    return [
+        "def _pack_fixed_str(text: str, capacity: int) -> bytes:",
+        '    """u32 byte length plus ``capacity`` reserved payload bytes."""',
+        '    raw = text.encode("utf-8")',
+        "    if len(raw) > capacity:",
+        "        raise ValueError(",
+        '            f"string is {len(raw)} bytes; FixedString[{capacity}] holds {capacity}"',
+        "        )",
+        '    return struct.pack("<I", len(raw)) + raw + bytes(capacity - len(raw))',
+    ]
+
+
+def _param_type(param: DataParamModel) -> str:
+    inner = _py_type(param.c_type)
+    if param.shape == "option":
+        return f"{inner} | None"
+    if param.shape == "string":
+        return "str"
+    return inner
+
+
+def _pack_expr(param: DataParamModel, name: str) -> str:
+    if param.shape == "option":
+        present = _pack_scalar(param.c_type, name)
+        width = param.elem_size
+        return (
+            f"(bytes([0]) + bytes({width})) if {name} is None else (bytes([1]) + {present})"
+        )
+    if param.shape == "string":
+        return f"_pack_fixed_str({name}, {param.length})"
+    return _pack_scalar(param.c_type, name)
+
+
+def _pack_scalar(c_type: str, name: str) -> str:
+    if c_type == "SolPubkey":
+        return f"bytes({name})"
+    if c_type == "uint128_t":
+        return f'{name}.to_bytes(16, "little")'
+    if c_type == "int128_t":
+        return f'{name}.to_bytes(16, "little", signed=True)'
+    if c_type == "bool":
+        return f'struct.pack("<B", 1 if {name} else 0)'
+    fmt = _PACK_FMT[c_type]
+    return f'struct.pack("{fmt}", {name})'
+
+
+def _instruction_fn(instr: InstructionModel) -> list[str]:
+    used = {"program_id"}
+    params: list[tuple[str, str]] = [("program_id", "Pubkey")]
+    account_names: list[tuple[str, InstructionArg]] = []
+    for account in instr.account_params:
+        py_name = _unique(account.name, used)
+        params.append((py_name, "Pubkey"))
+        account_names.append((py_name, InstructionArg(account.kind, account.writable)))
+    data_names: list[tuple[str, DataParamModel]] = []
+    for data in instr.data_params:
+        py_name = _unique(data.name, used)
+        params.append((py_name, _param_type(data)))
+        data_names.append((py_name, data))
+
+    lines = _signature(instr.snake_name, params, "Instruction")
+    lines.append(f'    """Build the `{instr.snake_name}` instruction."""')
+    pays = any(account.is_init for account in instr.account_params)
+    needs_system = pays or instr.needs_system
+    if data_names:
+        parts = [f"_DISC_{instr.snake_name.upper()}"]
+        parts.extend(_pack_expr(param, name) for name, param in data_names)
+        lines.append("    data = b''.join((")
+        for part in parts:
+            lines.append(f"        {part},")
+        lines.append("    ))")
+    else:
+        lines.append(f"    data = _DISC_{instr.snake_name.upper()}")
+    lines.append("    accounts = [")
+    for name, arg in account_names:
+        signer = arg.kind == "signer"
+        writable = arg.writable or (signer and pays)
+        lines.append(
+            f"        AccountMeta({name}, is_signer={signer}, is_writable={writable}),"
+        )
+    if needs_system:
+        lines.append(
+            "        AccountMeta(_SYSTEM_PROGRAM, is_signer=False, is_writable=False),"
+        )
+    if instr.needs_token or instr.needs_ata:
+        lines.append(
+            "        AccountMeta(_TOKEN_PROGRAM, is_signer=False, is_writable=False),"
+        )
+    if instr.needs_token2022:
+        lines.append(
+            "        AccountMeta(_TOKEN_2022_PROGRAM, is_signer=False, is_writable=False),"
+        )
+    if instr.needs_ata:
+        lines.append(
+            "        AccountMeta(_ATA_PROGRAM, is_signer=False, is_writable=False),"
+        )
+    lines.append("    ]")
+    lines.append("    return Instruction(program_id=program_id, data=data, accounts=accounts)")
+    return lines
+
+
+class InstructionArg:
+    """Signer/writable facts the client needs for one account parameter."""
+
+    def __init__(self, kind: str, writable: bool) -> None:
+        self.kind = kind
+        self.writable = writable
+
+
+def _decode_fn(
+    account: AccountModel,
+    structs: dict[str, StructModel],
+    payloads: dict[str, str] | None = None,
+) -> list[str]:
+    fn = f"decode_{snake_case(account.name)}"
+    doc = f"Decode a `{account.name}` account (discriminator plus fields)."
+    if any(field.shape == "option" for field in account.fields):
+        doc += " Option fields are (tag, value); the value bytes are always reserved."
+    if any(field.shape == "string" for field in account.fields):
+        doc += " Fixed strings are UTF-8 text; the reserved bytes stay on the account."
+    if any(field.shape == "struct" for field in account.fields):
+        doc += " Struct fields are dicts packed inline, with no extra discriminator."
+    lines = [
+        f"def {fn}(data: bytes) -> dict[str, object]:",
+        f'    """{doc}"""',
+        f"    if len(data) < {account.size}:",
+        f'        raise ValueError("{account.name} account is shorter than {account.size} bytes")',
+    ]
+    for field in account.fields:
+        lines.extend(
+            _decode_field(field, ACCOUNT_DATA_OFFSET + field.offset, structs, payloads)
+        )
+    lines.append("    return {")
+    for field in account.fields:
+        lines.append(f'        "{field.name}": {field.name},')
+    lines.append("    }")
+    return lines
+
+
+def _struct_decode_fn(
+    struct: StructModel,
+    structs: dict[str, StructModel],
+    payloads: dict[str, str] | None = None,
+) -> list[str]:
+    """Decode one ``@struct`` at a caller-supplied absolute offset."""
+    lines = [
+        f"def _decode_struct_{struct.name}(data: bytes, base: int) -> dict[str, object]:",
+        f'    """Decode `{struct.name}` at `base` (fixed width, no discriminator)."""',
+    ]
+    for field in struct.fields:
+        lines.extend(_decode_field(field, f"base + {field.offset}", structs, payloads))
+    lines.append("    return {")
+    for field in struct.fields:
+        lines.append(f'        "{field.name}": {field.name},')
+    lines.append("    }")
+    return lines
+
+
+def _decode_field(
+    field: FieldModel,
+    base: int | str,
+    structs: dict[str, StructModel],
+    payloads: dict[str, str] | None = None,
+) -> list[str]:
+    payload_c = "" if payloads is None else payloads.get(field.defined, "")
+    if field.shape == "scalar" and payload_c and field.size > 1:
+        at = base + 1 if isinstance(base, int) else f"({base}) + 1"
+        fmt = _UNPACK_FMT[payload_c]
+        return [
+            f"    {field.name}_tag = data[{base}]",
+            f'    ({field.name}_value,) = struct.unpack_from("{fmt}", data, {at})',
+            f"    {field.name} = ({field.name}_tag, {field.name}_value)",
+        ]
+    if field.shape == "struct":
+        if field.defined not in structs:
+            return [
+                f'    raise ValueError("unknown struct {field.defined}")',
+            ]
+        return [f"    {field.name} = _decode_struct_{field.defined}(data, {base})"]
+    if field.shape == "array":
+        return _decode_array(field, base)
+    if field.shape == "vec":
+        return _decode_vec(field, base)
+    if field.shape == "option":
+        return _decode_option(field, base)
+    if field.shape == "string":
+        return _decode_string(field, base)
+    return [_read_scalar(field.name, field.c_type, base, field.elem_size or field.size)]
+
+
+def _plus(base: int | str, delta: int) -> int | str:
+    """Absolute offset ``base + delta``. An int stays an int so account lines do not change."""
+    if isinstance(base, int):
+        return base + delta
+    if delta == 0:
+        return base
+    return f"({base}) + {delta}"
+
+
+def _decode_string(field: FieldModel, base: int | str) -> list[str]:
+    """``u32`` length at ``base``, reserved bytes at ``base + 4``. Length past ``N`` raises."""
+    length = f"_str_{field.name}_len"
+    raw = f"_str_{field.name}_raw"
+    payload = _plus(base, 4)
+    return [
+        f'    ({length},) = struct.unpack_from("<I", data, {base})',
+        f"    if {length} > {field.length}:",
+        f'        raise ValueError("{field.name} is longer than {field.length} bytes")',
+        f"    {raw} = data[{payload}:{payload} + {length}]",
+        f'    {field.name} = {raw}.decode("utf-8")',
+    ]
+
+
+def _decode_option(field: FieldModel, base: int | str) -> list[str]:
+    """Tag byte at ``base``, reserved inner value at ``base + 1``."""
+    tag = f"_opt_{field.name}_tag"
+    value = f"_opt_{field.name}_value"
+    lines = [f"    {tag} = data[{base}]"]
+    lines.append(_read_scalar(value, field.c_type, _plus(base, 1), field.elem_size))
+    lines.append(f"    {field.name} = ({tag}, {value})")
+    return lines
+
+
+def _decode_vec(field: FieldModel, base: int | str) -> list[str]:
+    """``u32`` length at ``base``, reserved elements at ``base + 4``."""
+    step = field.elem_size or 1
+    payload = _plus(base, 4)
+    length = f"_vec_{field.name}_len"
+    py = _py_type(field.c_type)
+    lines = [
+        f'    ({length},) = struct.unpack_from("<I", data, {base})',
+        f"    if {length} > {field.length}:",
+        f'        raise ValueError("{field.name} is longer than {field.length} elements")',
+        f"    {field.name}: list[{py}] = []",
+        f"    for _index in range({length}):",
+        f"        _at = {payload} + _index * {step}",
+    ]
+    elem = field.c_type
+    if elem == "SolPubkey":
+        lines.append(f"        {field.name}.append(Pubkey.from_bytes(data[_at:_at + 32]))")
+        return lines
+    if elem in ("uint128_t", "int128_t"):
+        signed = "True" if elem == "int128_t" else "False"
+        lines.append(
+            f"        {field.name}.append("
+            f'int.from_bytes(data[_at:_at + 16], "little", signed={signed}))'
+        )
+        return lines
+    if elem == "bool":
+        lines.append(f"        {field.name}.append(bool(data[_at]))")
+        return lines
+    fmt = _UNPACK_FMT[elem]
+    lines.append(f'        (_item,) = struct.unpack_from("{fmt}", data, _at)')
+    lines.append(f"        {field.name}.append(_item)")
+    return lines
+
+
+def _decode_array(field: FieldModel, base: int | str) -> list[str]:
+    elem = field.c_type
+    py = _py_type(elem)
+    step = field.elem_size or field.size // field.length
+    at = f"{base} + _index * {step}" if isinstance(base, int) else f"({base}) + _index * {step}"
+    lines = [f"    {field.name}: list[{py}] = []"]
+    lines.append(f"    for _index in range({field.length}):")
+    lines.append(f"        _at = {at}")
+    if elem == "SolPubkey":
+        lines.append(f"        {field.name}.append(Pubkey.from_bytes(data[_at:_at + 32]))")
+        return lines
+    if elem in ("uint128_t", "int128_t"):
+        signed = "True" if elem == "int128_t" else "False"
+        lines.append(
+            f"        {field.name}.append("
+            f'int.from_bytes(data[_at:_at + 16], "little", signed={signed}))'
+        )
+        return lines
+    if elem == "bool":
+        lines.append(f"        {field.name}.append(bool(data[_at]))")
+        return lines
+    fmt = _UNPACK_FMT[elem]
+    lines.append(f'        (_item,) = struct.unpack_from("{fmt}", data, _at)')
+    lines.append(f"        {field.name}.append(_item)")
+    return lines
+
+
+def _read_scalar(name: str, c_type: str, offset: int | str, size: int) -> str:
+    if c_type == "SolPubkey":
+        end = _plus(offset, 32)
+        return f"    {name} = Pubkey.from_bytes(data[{offset}:{end}])"
+    if c_type == "uint128_t":
+        end = _plus(offset, 16)
+        return f'    {name} = int.from_bytes(data[{offset}:{end}], "little")'
+    if c_type == "int128_t":
+        end = _plus(offset, 16)
+        return f'    {name} = int.from_bytes(data[{offset}:{end}], "little", signed=True)'
+    if c_type == "bool":
+        return f"    {name} = bool(data[{offset}])"
+    if size == 1 and c_type == "uint8_t":
+        return f"    {name} = data[{offset}]"
+    fmt = _UNPACK_FMT[c_type]
+    return f'    ({name},) = struct.unpack_from("{fmt}", data, {offset})'
+
+
+def _event_decoder(events: list[EventModel]) -> list[str]:
+    """``decode_event`` reads one ``sol_log_data`` slice (disc plus payload)."""
+    lines = ["# Event discriminators: sha256('event:' + name)[:8]."]
+    for event in events:
+        lines.append(f"_EVENT_{event.name.upper()} = {_bytes_expr(event.discriminator)}")
+    lines.append("")
+    lines.extend(
+        [
+            "def decode_event(data: bytes) -> dict[str, object]:",
+            '    """Decode one event log (8-byte discriminator, then packed fields)."""',
+            "    if len(data) < 8:",
+            '        raise ValueError("event is shorter than 8 bytes")',
+            "    disc = data[:8]",
+        ]
+    )
+    for event in events:
+        lines.append(f"    if disc == _EVENT_{event.name.upper()}:")
+        for field in event.fields:
+            raw = _read_scalar(field.name, field.c_type, 8 + field.offset, field.size)
+            lines.append("    " + raw)
+        lines.append("        return {")
+        lines.append(f'            "name": {event.name!r},')
+        lines.append('            "fields": {')
+        for field in event.fields:
+            lines.append(f'                "{field.name}": {field.name},')
+        lines.append("            },")
+        lines.append("        }")
+    lines.append('    raise ValueError("unknown event")')
+    return lines
+
+
+def _signature(name: str, params: list[tuple[str, str]], result: str) -> list[str]:
+    parts = [f"{param}: {ann}" for param, ann in params]
+    one = f"def {name}(" + ", ".join(parts) + f") -> {result}:"
+    if len(one) <= 100:
+        return [one]
+    lines = [f"def {name}("]
+    for part in parts:
+        lines.append(f"    {part},")
+    lines.append(f") -> {result}:")
+    return lines
+
+
+def _unique(name: str, used: set[str]) -> str:
+    candidate = name
+    while candidate in used:
+        candidate += "_"
+    used.add(candidate)
+    return candidate