Convert cffi type to numpy dtype

ctype_to_dtype.py

import numpy as np
import six

_FLOAT_TYPES = set(['float', 'double', 'long double'])

def _sub_overrides(overrides, prefix):
    out = {}
    for (key, value) in six.iteritems(overrides):
        if key.startswith(prefix):
            out[key[len(prefix):]] = value
    return out

def ctype_to_dtype(ffi, ctype, overrides=None):
    """Convert a CFFI ctype representing a struct to a numpy dtype.

    This does not necessarily handle all possible cases, but should correctly
    account for things like padding.

    Parameters
    ----------
    ffi
        cffi object imported from the library module
    ctype : `CType`
        Type object created from CFFI
    overrides : dict, optional
        Map elements of the type to specified numpy types. The keys are
        strings. To specify an element of a structure, use ``.name``. To
        specify the elements of an array type, use ``[]``. These strings can
        be concatenated (with no whitespace) to select sub-elements.
    """
    if overrides is None:
        overrides = {}
    try:
        return overrides['']
    except KeyError:
        pass

    if ctype.kind == 'primitive':
        if ctype.cname in _FLOAT_TYPES:
            return np.dtype('f' + str(ffi.sizeof(ctype)))
        elif ctype.cname == 'char':
            return np.dtype('c')
        elif ctype.cname == '_Bool':
            return np.dtype(np.bool_)
        else:
            test = int(ffi.cast(ctype, -1))
            if test == -1:
                return np.dtype('i' + str(ffi.sizeof(ctype)))
            else:
                return np.dtype('u' + str(ffi.sizeof(ctype)))
    elif ctype.kind == 'struct':
        names = []
        formats = []
        offsets = []
        for field in ctype.fields:
            if field[1].bitsize != -1:
                raise ValueError('bitfields are not supported')
            names.append(field[0])
            sub_overrides = _sub_overrides(overrides, '.' + field[0])
            formats.append(ctype_to_dtype(ffi, field[1].type, sub_overrides))
            offsets.append(field[1].offset)
        return np.dtype(dict(names=names, formats=formats, offsets=offsets, itemsize=ffi.sizeof(ctype)))
    elif ctype.kind == 'array':
        shape = []
        prefix = ''
        while ctype.kind == 'array' and prefix not in overrides:
            shape.append(ctype.length)
            ctype = ctype.item
            prefix += '[]'
        sub_overrides = _sub_overrides(overrides, prefix)
        return np.dtype((ctype_to_dtype(ffi, ctype, sub_overrides), tuple(shape)))
    else:
        raise ValueError('Unhandled kind {}'.format(ctype.kind))