Function test_build_inline_with_docstrings

tests/python/test_build.py:199–275

Test building functions with documentation using the functions dict.

()

def test_build_inline_with_docstrings() -> None:
    """Test building functions with documentation using the functions dict."""
    # Keep module alive until all returned objects are destroyed
    add_docstring = (
        "Add two integers and return the sum.\n"
        "\n"
        "Parameters\n"
        "----------\n"
        "a : int\n"
        "    First integer\n"
        "b : int\n"
        "    Second integer\n"
        "\n"
        "Returns\n"
        "-------\n"
        "result : int\n"
        "    Sum of a and b"
    )

    divide_docstring = "Divides two floats. Returns a/b."

    mod: Module = tvm_ffi.cpp.load_inline(
        name="test_docs",
        cpp_sources=r"""
            int add(int a, int b) {
                return a + b;
            }

            int subtract(int a, int b) {
                return a - b;
            }

            float divide(float a, float b) {
                TVM_FFI_ICHECK(b != 0.0f) << "Division by zero";
                return a / b;
            }
        """,
        functions={
            "add": add_docstring,
            "subtract": "",  # No documentation
            "divide": divide_docstring,
        },
        extra_cflags=["-DTVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1"],
    )

    # Test add function with full documentation
    assert mod.add(10, 5) == 15
    metadata = mod.get_function_metadata("add")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[int, int], int]"

    doc = mod.get_function_doc("add")
    assert doc is not None, "add should have documentation"
    assert doc == add_docstring

    # Test subtract function without documentation
    assert mod.subtract(10, 5) == 5