Vyper interface generator

vyper_interface.vy

"""
Tools for working with ABI specs and Vyper interface source code

```python
>>> from package import generate_interface
>>> generate_interface([{
        "type": "function",
        "name": "symbol",
        "stateMutability": "view",
        "inputs": [],
        "outputs": [{"type": "string"}],
    }], "ERC20Symbol")
'interface ERC20Symbol:\n    def symbol() -> string: view\n'
```
"""

from typing import Any, Dict, List, Optional, Union

from ethpm_types import (
    ABI,
    MethodABI,
)
from ethpm_types.abi import ABIType

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 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 abi_to_type(iface: Dict[str, Any]) -> Optional[ABI]:
    """Convert a dict JSON-like interface to an ethpm-types ABI type"""
    if iface["type"] == "function":
        return MethodABI.model_validate(iface)
    return None


def generate_interface(abi: List[Union[Dict[str, Any], 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, dict):
            iface = abi_to_type(iface)
            if iface is None:
                continue

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

    return source


__all__ = ["generate_interface"]