Function test_function_subclass

tests/python/test_function.py:281–312

()

def test_function_subclass() -> None:
    class JitFunction:
        def __init__(self, metadata: Any) -> None:
            self.metadata = metadata

    class MyFunction(tvm_ffi.Function, JitFunction):
        def __init__(self, metadata: Any) -> None:
            # Explicitly initialize the mixin. `super()` is not used because `tvm_ffi.Function`
            # is an extension type without a standard `__init__`.
            JitFunction.__init__(self, metadata)

        # When subclassing a Cython cdef class and overriding `__init__`,
        # special methods like `__call__` may not be inherited automatically.
        # This explicit assignment ensures the subclass remains callable.
        __call__ = tvm_ffi.Function.__call__

    f = tvm_ffi.convert(lambda x: x)
    assert isinstance(f, tvm_ffi.Function)
    f_sub = MyFunction(128)
    # move handle from f to f_sub an existing function
    f_sub.__move_handle_from__(f)
    assert isinstance(f_sub, MyFunction)
    assert isinstance(f_sub, JitFunction)
    assert f_sub.metadata == 128

    y: int = f_sub(2)
    assert y == 2
    echo = tvm_ffi.get_global_func("testing.echo")
    fechoed = echo(f_sub)
    assert isinstance(fechoed, tvm_ffi.Function)
    assert fechoed.__chandle__() == f_sub.__chandle__()
    assert fechoed(10) == 10


def test_function_with_opaque_ptr_protocol() -> None: