Function register_error

python/tvm_ffi/error.py:205–256

Register an error class so it can be recognized by the ffi error handler. Parameters ---------- name_or_cls The name of the error class. cls The class to register. Returns ------- fregister Register function if f is not specified. Examples

(
    name_or_cls: str | type | None = None,
    cls: type | None = None,
)

def register_error(
    name_or_cls: str | type | None = None,
    cls: type | None = None,
) -> Any:
    """Register an error class so it can be recognized by the ffi error handler.

    Parameters
    ----------
    name_or_cls
        The name of the error class.

    cls
        The class to register.

    Returns
    -------
    fregister
        Register function if f is not specified.

    Examples
    --------
    .. code-block:: python

        import tvm_ffi


        # Register a custom Python exception so tvm_ffi.Error maps to it
        @tvm_ffi.error.register_error
        class MyError(RuntimeError):
            pass


        # Convert a Python exception to an FFI Error and back
        ffi_err = tvm_ffi.convert(MyError("boom"))
        py_err = ffi_err.py_error()
        assert isinstance(py_err, MyError)

    """
    if isinstance(name_or_cls, type):
        cls = name_or_cls
        name_or_cls = cls.__name__

    def register(mycls: type) -> type:
        """Register the error class name with the FFI core."""
        err_name = name_or_cls if isinstance(name_or_cls, str) else mycls.__name__
        core.ERROR_NAME_TO_TYPE[err_name] = mycls
        core.ERROR_TYPE_TO_NAME[mycls] = err_name
        return mycls

    if cls is None:
        return register
    return register(cls)


register_error("RuntimeError", RuntimeError)