Vyper inline interface generator

generate_interface.py

"""
Generate an inline interface from a given Vyper source file.

NOTES
- We can probably use the same code to handle interfaces and contracts for 
    building ABIs, rather than having to use the compiler for the latter.
"""

from pathlib import Path
from typing import List

import click
from ethpm_types import (
    ABI,
    MethodABI,
)
from ethpm_types.abi import ABIType
from vyper.ast import parse_to_ast_with_settings
from vyper.ast.nodes import FunctionDef, Module, Name, Pass, Subscript


ERC20_LIMITED_FILE = Path(__file__).parent.parent / "contracts/interface/ERC20Limited.vy"
DEFAULT_VYPER_MUTABILITY = "nonpayable"
DECORATOR_MUTABILITY = {
    "pure",  # Function does not read contract state or environment variables
    "view",  # Function does not alter contract state
    "payable",  # Function is able to receive Ether and may alter state
    "nonpayable",  # Function may alter sate
}
INDENT_SPACES = 4
INDENT = " " * INDENT_SPACES


def indent_line(line: str, level=1) -> str:
    """Indent a source line of code"""
    return f"{INDENT * level}{line}"


def is_interface(module: Module):
    """
    Check if the module is an interface.  All functions body must be a pass.
    """
    for child in module.get_children():
        if isinstance(child, FunctionDef):
            func_body = child.get("body")
            if func_body is not None and not isinstance(func_body, Pass):
                return False
    return True


def funcdef_decorators(funcdef: FunctionDef) -> List[str]:
    return funcdef.get("decorator_list") or []


def funcdef_inputs(funcdef: FunctionDef) -> List[ABIType]:
    """Get a FunctionDef's defined input args"""
    args = funcdef.get("args")
    # TODO: Does Vyper allow complex input types, like structs and arrays?
    return (
        [ABIType.model_validate({"name": arg.arg, "type": arg.annotation.id}) for arg in args.args]
        if args
        else []
    )


def funcdef_outputs(funcdef: FunctionDef) -> List[ABIType]:
    """Get a FunctionDef's outputs, or return values"""
    returns = funcdef.get("returns")

    if not returns:
        return []

    if isinstance(returns, Name):
        # TODO: Structs fall in here. I think they're supposed to be a tuple of types in the ABI.
        #       Need to dig into that more.
        return [ABIType.model_validate({"type": returns.id})]

    elif isinstance(returns, Subscript):
        # An array type
        length = returns.slice.value.value
        array_type = returns.value.id
        # TOOD: Is this an acurrate way to define a fixed length array for ABI?
        return [ABIType.model_validate({"type": f"{array_type}[{length}]"})]

    raise NotImplementedError(f"Unhandled return type {type(returns)}")


def funcdef_state_mutability(funcdef: FunctionDef) -> str:
    """Get a FunctionDef's declared state mutability"""
    for decorator in funcdef_decorators(funcdef):
        if decorator in DECORATOR_MUTABILITY:
            return decorator
    return DEFAULT_VYPER_MUTABILITY


def funcdef_is_external(funcdef: FunctionDef) -> bool:
    """Check if a FunctionDef is declared external"""
    for decorator in funcdef_decorators(funcdef):
        if decorator == "external":
            return True
    return False


def funcdef_to_abi(func: FunctionDef) -> ABI:
    """Return a MethodABI instance for a Vyper FunctionDef"""
    return MethodABI.model_validate(
        {
            "name": func.get("name"),
            "inputs": funcdef_inputs(func),
            "outputs": funcdef_outputs(func),
            "stateMutability": funcdef_state_mutability(func),
        }
    )


def generate_abi(module: Module) -> List[ABI]:
    """
    Generate an ABI from a Vyper AST Module instance.
    """
    abi = []
    for child in module.get_children():
        if isinstance(child, FunctionDef):
            abi.append(funcdef_to_abi(child))
    return abi


def generate_inputs(inputs: List[ABIType]) -> str:
    """Generate the source code input args from ABI inputs"""
    return ", ".join(f"{i.name}: {i.type}" for i in inputs)


def generate_method(abi: MethodABI) -> str:
    """Generate Vyper interface method definition"""
    inputs = generate_inputs(abi.inputs)
    return_maybe = f" -> {abi.outputs[0].type}" if abi.outputs else ""
    return f"def {abi.name}({inputs}){return_maybe}: {abi.stateMutability}\n"


def generate_interface(abi: List[ABI], iface_name) -> str:
    """
    Generate a Vyper interface source code from an ABI spec
    Args:
        abi (List[Union[Dict[str, Any], ABI]]): An ABI spec for a contract
        iface_name (str): The name of the interface
    Returns:
        ``str`` Vyper source code for the interface
    """
    source = f"interface {iface_name}:\n"

    for iface in abi:
        if isinstance(iface, MethodABI):
            source += indent_line(generate_method(iface))

    return source


@click.command()
@click.argument("interface_name")  # The name of the interface to generate
@click.argument("vyper_file", type=Path)  # The file to generate the interface from
def main(interface_name: str, vyper_file: Path):
    source = vyper_file.read_text()
    settings, ast = parse_to_ast_with_settings(source)
    abi = generate_abi(ast)
    print(generate_interface(abi, interface_name))


if __name__ == "__main__":
    main()