enthought · prabhuramachandran · Jun 28, 2018 · Jun 28, 2018
diff --git a/tvtk/array_handler.py b/tvtk/array_handler.py
@@ -25,7 +25,12 @@
 import numpy
 
 # Enthought library imports.
-from tvtk.array_ext import set_id_type_array
+try:
+    from tvtk.array_ext import set_id_type_array
+    HAS_ARRAY_EXT = True
+except ImportError:
+    HAS_ARRAY_EXT = False
+
 from tvtk.common import is_old_pipeline
 
 # Useful constants for VTK arrays.
@@ -48,10 +53,47 @@
 if sys.version_info[0] > 2:
     unicode = str
 
+
 def getbuffer(array):
     return getattr(numpy, 'getbuffer', memoryview)(array)
 
 
+def set_id_type_array_py(id_array, out_array):
+    """Given a 2D Int array (`id_array`), and a contiguous 1D numarray array
+    (`out_array`) having the correct size, this function sets the data from
+    `id_array` into `out_array` so that it can be used in place of a
+    `vtkIdTypeArray` in order to set the cells of a `vtkCellArray`.
+
+    Note that if `shape = id_array.shape` then `size(out_array) ==
+    shape[0]*(shape[1] + 1)` should be true. If not you'll get an
+    `AssertionError`.
+
+    `id_array` need not be contiguous but `out_array` must be.
+
+    """
+    assert numpy.issubdtype(id_array.dtype, numpy.signedinteger)
+    assert out_array.flags.contiguous == 1, \
+        "out_array must be contiguous."
+    shp = id_array.shape
+    assert len(shp) == 2, "id_array must be a two dimensional array."
+    sz = numpy.size(out_array)
+    e_sz = shp[0]*(shp[1]+1)
+    assert sz == e_sz, \
+        "out_array size is incorrect, expected: %s, given: %s" % (e_sz, sz)
+
+    # numpy.insert does the trick for us albeit slower than our Cython
+    # implementation (by as much as 4x).
+    res = numpy.insert(id_array, 0, shp[1], axis=1)
+    if len(out_array.shape) > 1:
+        out_array[:] = res
+    else:
+        out_array[:] = res.ravel()
+
+
+if not HAS_ARRAY_EXT:
+    set_id_type_array = set_id_type_array_py
+
+
 ######################################################################
 # The array cache.
 ######################################################################

diff --git a/tvtk/setup.py b/tvtk/setup.py
@@ -4,19 +4,37 @@
 #
 from __future__ import print_function
 
-import os, sys
+import os
+import sys
+
+
+def can_compile_extensions():
+    from distutils.dist import Distribution
+    sargs = {'script_name': None, 'script_args': ["--build-ext"]}
+    d = Distribution(sargs)
+    cfg = d.get_command_obj('config')
+    cfg.dump_source = 0
+    cfg.noisy = 0
+    cfg.finalize_options()
+    build_ext = d.get_command_obj('build_ext')
+    build_ext.finalize_options()
+    return cfg.try_compile(
+        'int main(void) {return 0;}',
+        headers=['Python.h'],
+        include_dirs=build_ext.include_dirs,
+        lang='c'
+    )
 
 
 def configuration(parent_package=None, top_path=None):
     from os.path import join
     from numpy.distutils.misc_util import Configuration
-    config = Configuration('tvtk',parent_package,top_path)
+    config = Configuration('tvtk', parent_package, top_path)
     config.set_options(ignore_setup_xxx_py=True,
                        assume_default_configuration=True,
                        delegate_options_to_subpackages=True,
                        quiet=True)
 
-
     config.add_subpackage('custom')
     config.add_subpackage('pipeline')
     config.add_subpackage('pyface')
@@ -37,15 +55,18 @@ def configuration(parent_package=None, top_path=None):
 
     config.add_subpackage('tests')
 
-    # Numpy support.
-    config.add_extension('array_ext',
-                         sources = [join('src','array_ext.c')],
-                         depends = [join('src','array_ext.pyx')],
-                         )
+    # Add any extensions.  These are optional.
+    if can_compile_extensions():
+        config.add_extension(
+            'array_ext',
+            sources=[join('src', 'array_ext.c')],
+            depends=[join('src', 'array_ext.pyx')],
+        )
 
     tvtk_classes_zip_depends = config.paths(
-        'code_gen.py','wrapper_gen.py', 'special_gen.py',
-        'tvtk_base.py', 'indenter.py', 'vtk_parser.py')
+        'code_gen.py', 'wrapper_gen.py', 'special_gen.py',
+        'tvtk_base.py', 'indenter.py', 'vtk_parser.py'
+    )
 
     return config
 

diff --git a/tvtk/tests/test_array_ext.py b/tvtk/tests/test_array_ext.py
@@ -9,34 +9,35 @@
 
 import numpy
 
-from tvtk.array_handler import ID_TYPE_CODE
+from tvtk.array_handler import ID_TYPE_CODE, set_id_type_array_py
 from tvtk.array_ext import set_id_type_array
 
+
 class TestArrayExt(unittest.TestCase):
-    def test_set_id_type_array(self):
+    def check(self, set_id_type_array):
         N = 5
-        a = numpy.zeros((N,4), ID_TYPE_CODE)
-        a[:,1] = 1
-        a[:,2] = 2
-        a[:,3] = 3
+        a = numpy.zeros((N, 4), ID_TYPE_CODE)
+        a[:, 1] = 1
+        a[:, 2] = 2
+        a[:, 3] = 3
 
         def diff_arr(x, y):
-            return numpy.sum(numpy.ravel(x) - numpy.ravel(y[:,1:]))
+            return numpy.sum(numpy.ravel(x) - numpy.ravel(y[:, 1:]))
 
         # Test contiguous arrays.
-        b = numpy.zeros((N,5), ID_TYPE_CODE)
+        b = numpy.zeros((N, 5), ID_TYPE_CODE)
         set_id_type_array(a, b)
         self.assertEqual(diff_arr(a, b), 0)
 
         # Test non-contiguous arrays.
-        b = numpy.zeros((N,3), ID_TYPE_CODE)
-        set_id_type_array(a[:,::2], b)
-        self.assertEqual(diff_arr(a[:,::2], b), 0)
+        b = numpy.zeros((N, 3), ID_TYPE_CODE)
+        set_id_type_array(a[:, ::2], b)
+        self.assertEqual(diff_arr(a[:, ::2], b), 0)
 
         # Test 1D array.
         b = numpy.zeros(N*5, ID_TYPE_CODE)
         set_id_type_array(a, b)
-        self.assertEqual(diff_arr(a, numpy.reshape(b, (N,5))), 0)
+        self.assertEqual(diff_arr(a, numpy.reshape(b, (N, 5))), 0)
 
         # Test assertions.
         d = a.astype('d')
@@ -47,7 +48,7 @@ def diff_arr(x, y):
         # B should b contiguous.
         b = numpy.zeros((N, 10), ID_TYPE_CODE)
         self.assertRaises(AssertionError, set_id_type_array,
-                          a, b[:,::2])
+                          a, b[:, ::2])
 
         self.assertRaises(AssertionError, set_id_type_array,
                           a[0], b)
@@ -63,7 +64,13 @@ def diff_arr(x, y):
 
         # This should work!
         set_id_type_array(a, b[:N*5])
-        self.assertEqual(diff_arr(a, numpy.reshape(b[:N*5], (N,5))), 0)
+        self.assertEqual(diff_arr(a, numpy.reshape(b[:N*5], (N, 5))), 0)
+
+    def test_set_id_type_array(self):
+        self.check(set_id_type_array)
+
+    def test_set_id_type_array_py(self):
+        self.check(set_id_type_array_py)
 
 
 if __name__ == "__main__":