py::buffer_info to_py_buffer()

in src/mlio-py/mlio/core/py_device_array.cc [68:139]


py::buffer_info to_py_buffer(Py_device_array &arr)
{
    auto size = static_cast<py::ssize_t>(arr.size());

    std::size_t item_size{};
    std::string fmt{};

    switch (arr.data_type()) {
    case Data_type::size:
        item_size = sizeof(std::size_t);
#if SIZE_MAX == UINT32_MAX
        fmt = "i";
#elif SIZE_MAX == UINT64_MAX
        fmt = "q";
#else
#error "Only 32-bit and 64-bit systems are supported."
#endif
        break;
    case Data_type::float16:
        item_size = sizeof(std::uint16_t);
        fmt = "e";
        break;
    case Data_type::float32:
        item_size = sizeof(float);
        fmt = "f";
        break;
    case Data_type::float64:
        item_size = sizeof(double);
        fmt = "d";
        break;
    case Data_type::int8:
        item_size = sizeof(std::int8_t);
        fmt = "b";
        break;
    case Data_type::int16:
        item_size = sizeof(std::int16_t);
        fmt = "h";
        break;
    case Data_type::int32:
        item_size = sizeof(std::int32_t);
        fmt = "i";
        break;
    case Data_type::int64:
        item_size = sizeof(std::int64_t);
        fmt = "q";
        break;
    case Data_type::uint8:
        item_size = sizeof(std::uint8_t);
        fmt = "B";
        break;
    case Data_type::uint16:
        item_size = sizeof(std::uint16_t);
        fmt = "H";
        break;
    case Data_type::uint32:
        item_size = sizeof(std::uint32_t);
        fmt = "I";
        break;
    case Data_type::uint64:
        item_size = sizeof(std::uint64_t);
        fmt = "Q";
        break;
    case Data_type::string:
        item_size = sizeof(PyObject *);
        fmt = "O";
        break;
    }

    auto is = static_cast<py::ssize_t>(item_size);

    return py::buffer_info(arr.data(), is, fmt, size);
}