Function _make_kernel_ast

python/tests/kernel/test_kernel_signature.py:80–93

Build an ``ast.Module`` for a stub kernel with the given type annotations. ``arg_types`` is a list of annotation strings (e.g. ``["int", "float"]``). ``return_type`` is an annotation string (e.g. ``"bool"``) or ``None``.

(arg_types=None, return_type=None)

def _make_kernel_ast(arg_types=None, return_type=None) -> ast.Module:
    """
    Build an ``ast.Module`` for a stub kernel with the given type annotations.

    ``arg_types`` is a list of annotation strings (e.g. ``["int", "float"]``).
    ``return_type`` is an annotation string (e.g. ``"bool"``) or ``None``.
    """
    params = ", ".join(f"a{i}: {t}" for i, t in enumerate(arg_types or []))
    ret = f" -> {return_type}" if return_type else ""
    src = textwrap.dedent(f"""\
        def kernel({params}){ret}:
            pass
    """)
    return ast.parse(src)


# -- test cases -------------------------------------------------------------