1 import numpy as N
   2 import ctypes as C
   3 
   4 
   5 
   6 ctypesDict = {'d' : C.c_double,
   7               'b' : C.c_char,
   8               'h' : C.c_short,
   9               'i' : C.c_int,
  10               'l' : C.c_long,
  11               'q' : C.c_longlong,
  12               'B' : C.c_ubyte,
  13               'H' : C.c_ushort,
  14               'I' : C.c_uint,
  15               'L' : C.c_ulong,
  16               'Q' : C.c_ulonglong}
  17               
  18 
  19 
  20 
  21 
  22 
  23 def c_ndarray(a, dtype = None, ndim = None, shape = None, requirements = None):
  24 
  25     """
  26     PURPOSE: Given an array, return a ctypes structure containing the
  27              arrays info (data, shape, strides, ndim). A check is made to ensure
  28              that the array has the specified dtype and requirements.
  29              
  30     INPUT: a: an array: something that is or can be converted to a numpy array
  31            dtype: the required dtype of the array, convert if it doesn't match
  32            ndim: integer: the required number of axes of the array
  33            shape: tuple of integers: required shape of the array
  34            requirements: list of requirements: (E)nsurearray, (F)ortran, (F)_contiguous,
  35                         (C)ontiguous, (C)_contiguous. Convert if it doesn't match.
  36                          
  37     OUTPUT: ctypes structure with the fields:
  38              . data: pointer to the data : the type is determined with the dtype of the
  39                      array, and with ctypesDict.
  40              . shape: pointer to long array : size of each of the dimensions
  41              . strides: pointer to long array : strides in elements (not bytes)
  42     """
  43     
  44     if not requirements:
  45         
  46         # Also allow derived classes of ndarray
  47         array = N.asanyarray(a, dtype=dtype)
  48     else:
  49     
  50         # Convert requirements to captial letter codes:
  51         # (ensurearray' -> 'E';  'aligned' -> 'A'
  52         #  'fortran', 'f_contiguous', 'f' -> 'F'
  53         #  'contiguous', 'c_contiguous', 'c' -> 'C')
  54     
  55         requirements = [x[0].upper() for x in requirements]
  56         subok = (0 if 'E' in requirements else 1)
  57     
  58         # Make from 'a' an ndarray with the specified dtype, but don't copy the
  59         # data (yet). This also ensures that the .flags attribute is present.
  60         
  61         array = N.array(a, dtype=dtype, copy=False, subok=subok)
  62 
  63         # See if copying all data is really necessary.
  64         # Note: 'A' = (A)ny = only (F) it is was already (F) 
  65         
  66         copychar = 'A' 
  67         if 'F' in requirements:
  68             copychar = 'F'
  69         elif 'C' in requirements:
  70             copychar = 'C'
  71         
  72         for req in requirements:
  73             if not array.flags[req]:
  74                 array = array.copy(copychar)
  75                 break
  76     
  77     
  78     # If required, check the number of axes and the shape of the array
  79     
  80     if ndim is not None:
  81         if array.ndim != ndim:
  82             raise TypeError, "Array has wrong number of axes"
  83                 
  84     if shape is not None:
  85         if array.shape != shape:
  86             raise TypeError, "Array has wrong shape"
  87         
  88     # Define a class that serves as interface of an ndarray to ctypes.
  89     # Part of the type depends on the array's dtype. 
  90     
  91     class ndarrayInterfaceToCtypes(C.Structure):
  92         pass
  93         
  94     typechar = array.dtype.char
  95 
  96     if typechar in ctypesDict:
  97         ndarrayInterfaceToCtypes._fields_ =                       \
  98                   [("data", C.POINTER(ctypesDict[typechar])),
  99                    ("shape" , C.POINTER(C.c_long)),
 100                    ("strides", C.POINTER(C.c_long))]
 101     else:
 102         raise TypeError, "dtype of input ndarray not supported"
 103 
 104     # Instantiate the interface class and attach the ndarray's internal info.
 105     # Ctypes does automatic conversion between (c_long * #) arrays and POINTER(c_long).
 106     
 107     ndarrayInterface = ndarrayInterfaceToCtypes()
 108     ndarrayInterface.data = array.ctypes.data_as(C.POINTER(ctypesDict[typechar]))
 109     ndarrayInterface.shape = (C.c_long * array.ndim)(*array.shape)
 110     ndarrayInterface.strides = (C.c_long * array.ndim)(*array.strides)
 111     for n in range(array.ndim):
 112         ndarrayInterface.strides[n] /= array.dtype.itemsize
 113         
 114     return ndarrayInterface