kork/tests/test_foreign_functions.py

from typing import Any, Literal, Mapping, Sequence, Union

from kork.ast import ExternFunctionDef, Param, ParamList
from kork.foreign_funcs import to_extern_func_def


# Do not type the function below. We're testing initialization of the retriever.
def foo():  # type: ignore
    pass


def bar(x: int) -> int:
    """Add one to x."""
    return x + 1


def zoo(x: Mapping[str, Any], q: Sequence[Union[Literal["a"], Literal["b"]]]) -> Any:
    """Add one to x."""
    pass


def test_to_extern_func_def() -> None:
    """Test conversion of a python function into an extern function definition."""
    assert to_extern_func_def(foo) == ExternFunctionDef(
        name="foo",
        params=ParamList(params=[]),
        # If the source function is un-typed, we assume
        # that the return type is Any.
        return_type="Any",
        implementation=foo,
        doc_string="",
    )

    assert to_extern_func_def(bar) == ExternFunctionDef(
        name="bar",
        params=ParamList(params=[Param(name="x", type_="int")]),
        return_type="int",
        implementation=bar,
        doc_string="Add one to x.",
    )

    assert to_extern_func_def(zoo) == ExternFunctionDef(
        name="zoo",
        params=ParamList(
            params=[
                Param(name="x", type_="Mapping[str, Any]"),
                Param(
                    name="q",
                    type_="Sequence[Union[Literal['a'], Literal['b']]]",
                ),
            ]
        ),
        return_type="Any",
        implementation=zoo,
        doc_string="Add one to x.",
    )