Merge pull request #10970 from eric-wieser/cut-down-ctypeslib

WIP: Remove fragile use of  __array_interface__ in ctypeslib.as_array

2 files changed, 96 insertions, 114 deletions

diff --git a/numpy/ctypeslib.py b/numpy/ctypeslib.py
index b8457c78b..9d71adbdb 100644
--- a/numpy/ctypeslib.py
+++ b/numpy/ctypeslib.py
@@ -319,120 +319,47 @@ def ndpointer(dtype=None, ndim=None, shape=None, flags=None):
     _pointer_type_cache[(dtype, shape, ndim, num)] = klass
     return klass
 
-if ctypes is not None:
-    ct = ctypes
-    ################################################################
-    # simple types
-
-    # maps the numpy typecodes like '<f8' to simple ctypes types like
-    # c_double. Filled in by prep_simple.
-    _typecodes = {}
-
-    def prep_simple(simple_type, dtype):
-        """Given a ctypes simple type, construct and attach an
-        __array_interface__ property to it if it does not yet have one.
-        """
-        try: simple_type.__array_interface__
-        except AttributeError: pass
-        else: return
-
-        typestr = _dtype(dtype).str
-        _typecodes[typestr] = simple_type
-
-        def __array_interface__(self):
-            return {'descr': [('', typestr)],
-                    '__ref': self,
-                    'strides': None,
-                    'shape': (),
-                    'version': 3,
-                    'typestr': typestr,
-                    'data': (ct.addressof(self), False),
-                    }
-
-        simple_type.__array_interface__ = property(__array_interface__)
 
+def _get_typecodes():
+    """ Return a dictionary mapping __array_interface__ formats to ctypes types """
+    ct = ctypes
     simple_types = [
-        ((ct.c_byte, ct.c_short, ct.c_int, ct.c_long, ct.c_longlong), "i"),
-        ((ct.c_ubyte, ct.c_ushort, ct.c_uint, ct.c_ulong, ct.c_ulonglong), "u"),
-        ((ct.c_float, ct.c_double), "f"),
+        ct.c_byte, ct.c_short, ct.c_int, ct.c_long, ct.c_longlong,
+        ct.c_ubyte, ct.c_ushort, ct.c_uint, ct.c_ulong, ct.c_ulonglong,
+        ct.c_float, ct.c_double,
     ]
 
-    # Prep that numerical ctypes types:
-    for types, code in simple_types:
-        for tp in types:
-            prep_simple(tp, "%c%d" % (code, ct.sizeof(tp)))
+    return {_dtype(ctype).str: ctype for ctype in simple_types}
 
-    ################################################################
-    # array types
 
-    _ARRAY_TYPE = type(ct.c_int * 1)
+def _ctype_ndarray(element_type, shape):
+    """ Create an ndarray of the given element type and shape """
+    for dim in shape[::-1]:
+        element_type = element_type * dim
+    return element_type
 
-    def prep_array(array_type):
-        """Given a ctypes array type, construct and attach an
-        __array_interface__ property to it if it does not yet have one.
-        """
-        try: array_type.__array_interface__
-        except AttributeError: pass
-        else: return
-
-        shape = []
-        ob = array_type
-        while type(ob) is _ARRAY_TYPE:
-            shape.append(ob._length_)
-            ob = ob._type_
-        shape = tuple(shape)
-        ai = ob().__array_interface__
-        descr = ai['descr']
-        typestr = ai['typestr']
-
-        def __array_interface__(self):
-            return {'descr': descr,
-                    '__ref': self,
-                    'strides': None,
-                    'shape': shape,
-                    'version': 3,
-                    'typestr': typestr,
-                    'data': (ct.addressof(self), False),
-                    }
-
-        array_type.__array_interface__ = property(__array_interface__)
-
-    def prep_pointer(pointer_obj, shape):
-        """Given a ctypes pointer object, construct and
-        attach an __array_interface__ property to it if it does not
-        yet have one.
-        """
-        try: pointer_obj.__array_interface__
-        except AttributeError: pass
-        else: return
-
-        contents = pointer_obj.contents
-        dtype = _dtype(type(contents))
-
-        inter = {'version': 3,
-                 'typestr': dtype.str,
-                 'data': (ct.addressof(contents), False),
-                 'shape': shape}
-
-        pointer_obj.__array_interface__ = inter
 
-    ################################################################
-    # public functions
+if ctypes is not None:
+    _typecodes = _get_typecodes()
 
     def as_array(obj, shape=None):
-        """Create a numpy array from a ctypes array or a ctypes POINTER.
+        """
+        Create a numpy array from a ctypes array or POINTER.
+
         The numpy array shares the memory with the ctypes object.
 
-        The size parameter must be given if converting from a ctypes POINTER.
-        The size parameter is ignored if converting from a ctypes array
+        The shape parameter must be given if converting from a ctypes POINTER.
+        The shape parameter is ignored if converting from a ctypes array
         """
-        tp = type(obj)
-        try: tp.__array_interface__
-        except AttributeError:
-            if hasattr(obj, 'contents'):
-                prep_pointer(obj, shape)
-            else:
-                prep_array(tp)
+        if isinstance(obj, ctypes._Pointer):
+            # convert pointers to an array of the desired shape
+            if shape is None:
+                raise TypeError(
+                    'as_array() requires a shape argument when called on a '
+                    'pointer')
+            p_arr_type = ctypes.POINTER(_ctype_ndarray(obj._type_, shape))
+            obj = ctypes.cast(obj, p_arr_type).contents
+
         return array(obj, copy=False)
 
     def as_ctypes(obj):
@@ -446,9 +373,7 @@ if ctypes is not None:
         addr, readonly = ai["data"]
         if readonly:
             raise TypeError("readonly arrays unsupported")
-        tp = _typecodes[ai["typestr"]]
-        for dim in ai["shape"][::-1]:
-            tp = tp * dim
+        tp = _ctype_ndarray(_typecodes[ai["typestr"]], ai["shape"])
         result = tp.from_address(addr)
         result.__keep = ai
         return result
diff --git a/numpy/tests/test_ctypeslib.py b/numpy/tests/test_ctypeslib.py
index 0f0d6dbc4..75ce9c8ca 100644
--- a/numpy/tests/test_ctypeslib.py
+++ b/numpy/tests/test_ctypeslib.py
@@ -4,9 +4,9 @@ import sys
 import pytest
 
 import numpy as np
-from numpy.ctypeslib import ndpointer, load_library
+from numpy.ctypeslib import ndpointer, load_library, as_array
 from numpy.distutils.misc_util import get_shared_lib_extension
-from numpy.testing import assert_, assert_raises
+from numpy.testing import assert_, assert_array_equal, assert_raises, assert_equal
 
 try:
     cdll = None
@@ -21,11 +21,12 @@ try:
 except ImportError:
     _HAS_CTYPE = False
 
+
+@pytest.mark.skipif(not _HAS_CTYPE,
+                    reason="ctypes not available in this python")
+@pytest.mark.skipif(sys.platform == 'cygwin',
+                    reason="Known to fail on cygwin")
 class TestLoadLibrary(object):
-    @pytest.mark.skipif(not _HAS_CTYPE,
-                        reason="ctypes not available in this python")
-    @pytest.mark.skipif(sys.platform == 'cygwin',
-                        reason="Known to fail on cygwin")
     def test_basic(self):
         try:
             # Should succeed
@@ -35,10 +36,6 @@ class TestLoadLibrary(object):
                    " (import error was: %s)" % str(e))
             print(msg)
 
-    @pytest.mark.skipif(not _HAS_CTYPE,
-                        reason="ctypes not available in this python")
-    @pytest.mark.skipif(sys.platform == 'cygwin',
-                        reason="Known to fail on cygwin")
     def test_basic2(self):
         # Regression for #801: load_library with a full library name
         # (including extension) does not work.
@@ -54,6 +51,7 @@ class TestLoadLibrary(object):
                    " (import error was: %s)" % str(e))
             print(msg)
 
+
 class TestNdpointer(object):
     def test_dtype(self):
         dt = np.intc
@@ -113,3 +111,62 @@ class TestNdpointer(object):
         a1 = ndpointer(dtype=np.float64)
         a2 = ndpointer(dtype=np.float64)
         assert_(a1 == a2)
+
+
+@pytest.mark.skipif(not _HAS_CTYPE,
+                    reason="ctypes not available on this python installation")
+class TestAsArray(object):
+    def test_array(self):
+        from ctypes import c_int
+
+        pair_t = c_int * 2
+        a = as_array(pair_t(1, 2))
+        assert_equal(a.shape, (2,))
+        assert_array_equal(a, np.array([1, 2]))
+        a = as_array((pair_t * 3)(pair_t(1, 2), pair_t(3, 4), pair_t(5, 6)))
+        assert_equal(a.shape, (3, 2))
+        assert_array_equal(a, np.array([[1, 2], [3, 4], [5, 6]]))
+
+    def test_pointer(self):
+        from ctypes import c_int, cast, POINTER
+
+        p = cast((c_int * 10)(*range(10)), POINTER(c_int))
+
+        a = as_array(p, shape=(10,))
+        assert_equal(a.shape, (10,))
+        assert_array_equal(a, np.arange(10))
+
+        a = as_array(p, shape=(2, 5))
+        assert_equal(a.shape, (2, 5))
+        assert_array_equal(a, np.arange(10).reshape((2, 5)))
+
+        # shape argument is required
+        assert_raises(TypeError, as_array, p)
+
+    def test_struct_array_pointer(self):
+        from ctypes import c_int16, Structure, pointer
+
+        class Struct(Structure):
+            _fields_ = [('a', c_int16)]
+
+        Struct3 = 3 * Struct
+
+        c_array = (2 * Struct3)(
+            Struct3(Struct(a=1), Struct(a=2), Struct(a=3)),
+            Struct3(Struct(a=4), Struct(a=5), Struct(a=6))
+        )
+
+        expected = np.array([
+            [(1,), (2,), (3,)],
+            [(4,), (5,), (6,)],
+        ], dtype=[('a', np.int16)])
+
+        def check(x):
+            assert_equal(x.dtype, expected.dtype)
+            assert_equal(x, expected)
+
+        # all of these should be equivalent
+        check(as_array(c_array))
+        check(as_array(pointer(c_array), shape=()))
+        check(as_array(pointer(c_array[0]), shape=(2,)))
+        check(as_array(pointer(c_array[0][0]), shape=(2, 3)))