encapsulated VTK Info in class

Author: oesteban

nipype/algorithms/mesh.py

@@ -44,33 +44,47 @@
         del os.environ['ETS_TOOLKIT']
 
 
-def no_tvtk():
-    global have_tvtk
-    return not have_tvtk
+class Info(object):
+    """ Handle VTK version information """
+    _vtk_version = None
+
+    @staticmethod
+    def vtk_version():
+        """ Get VTK version """
+        if not Info.no_tvtk():
+            return None
+
+        if Info._vtk_version is None:
+            try:
+                from tvtk.tvtk_classes.vtk_version import vtk_build_version
+                vsplits = vtk_build_version.split('.')
+                Info._vtk_version = tuple([int(vsplits[0]), int(vsplits[1])] + vsplits[2:])
+            except ImportError:
+                iflogger.warning(
+                    'VTK version-major inspection using tvtk failed, assuming VTK == 4.0.')
+                Info._vtk_version = (4, 0)
+
+        return Info._vtk_version
+
+    @staticmethod
+    def no_tvtk():
+        global have_tvtk
+        return not have_tvtk
 
 
 class TVTKBaseInterface(BaseInterface):
 
     """ A base class for interfaces using VTK """
 
     _redirect_x = True
-    _vtk_version = (4, 0, 0)
 
     def __init__(self, **inputs):
-        if no_tvtk():
+        if Info.no_tvtk():
             raise ImportError('This interface requires tvtk to run.')
-
-        try:
-            from tvtk.tvtk_classes.vtk_version import vtk_build_version
-            vsplits = vtk_build_version.split('.')
-            self._vtk_version = tuple([int(vsplits[0]), int(vsplits[1])] + vsplits[2:])
-        except ImportError:
-            iflogger.warning(
-                'VTK version-major inspection using tvtk failed, assuming VTK == 4.0.')
         super(TVTKBaseInterface, self).__init__(**inputs)
 
-    def version(self):
-        return self._vtk_version
+    def vtk_version(self):
+        return Info.vtk_version()
 
 
 class WarpPointsInputSpec(BaseInterfaceInputSpec):
@@ -160,7 +174,7 @@ def _run_interface(self, runtime):
         newpoints = [p + d for p, d in zip(points, disps)]
         mesh.points = newpoints
         w = tvtk.PolyDataWriter()
-        if self.version()[0] < 6:
+        if self.vtk_version()[0] < 6:
             w.input = mesh
         else:
             w.set_input_data_object(mesh)
@@ -292,7 +306,7 @@ def _run_interface(self, runtime):
         writer = tvtk.PolyDataWriter(
             file_name=op.abspath(self.inputs.out_warp))
 
-        if self.version()[0] < 6:
+        if self.vtk_version()[0] < 6:
             writer.input = out_mesh
         else:
             writer.set_input_data_object(out_mesh)
@@ -413,7 +427,7 @@ def _run_interface(self, runtime):
         vtk1.point_data.vectors = warping
         writer = tvtk.PolyDataWriter(
             file_name=op.abspath(self.inputs.out_warp))
-        if self.version()[0] < 6:
+        if self.vtk_version()[0] < 6:
             writer.input = vtk1
         else:
             writer.set_input_data_object(vtk1)
@@ -424,7 +438,7 @@ def _run_interface(self, runtime):
         writer = tvtk.PolyDataWriter(
             file_name=op.abspath(self.inputs.out_file))
 
-        if self.version()[0] < 6:
+        if self.vtk_version()[0] < 6:
             writer.input = vtk1
         else:
             writer.set_input_data_object(vtk1)

nipype/algorithms/tests/test_mesh_ops.py

@@ -45,7 +45,7 @@ def test_trans_distances():
     curdir = os.getcwd()
     os.chdir(tempdir)
 
-    if m.no_tvtk():
+    if m.Info.no_tvtk():
         yield assert_raises, ImportError, m.ComputeMeshWarp
     else:
         from nipype.algorithms.mesh import tvtk
@@ -60,7 +60,12 @@ def test_trans_distances():
         vtk1.points = np.array(vtk1.points) + inc
 
         writer = tvtk.PolyDataWriter(file_name=warped_surf)
-        writer.set_input_data(vtk1)
+
+        if m.Info.vtk_version() < 6:
+            writer.set_input(vtk1)
+        else:
+            writer.set_input_data_object(vtk1)
+
         writer.write()
 
         dist = m.ComputeMeshWarp()
@@ -82,7 +87,7 @@ def test_warppoints():
     curdir = os.getcwd()
     os.chdir(tempdir)
 
-    if m.no_tvtk():
+    if m.Info.no_tvtk():
         yield assert_raises, ImportError, m.WarpPoints
 
     # TODO: include regression tests for when tvtk is installed
@@ -96,7 +101,7 @@ def test_meshwarpmaths():
     curdir = os.getcwd()
     os.chdir(tempdir)
 
-    if m.no_tvtk():
+    if m.Info.no_tvtk():
         yield assert_raises, ImportError, m.MeshWarpMaths
 
     # TODO: include regression tests for when tvtk is installed

nipype/interfaces/fsl/utils.py

@@ -34,6 +34,8 @@
 from ...utils.filemanip import (load_json, save_json, split_filename,
                                 fname_presuffix, copyfile)
 
+from ...algorithms.mesh import Info as VTKInfo
+
 warn = warnings.warn
 
 
@@ -1872,8 +1874,8 @@ def __init__(self, command=None, **inputs):
     def _format_arg(self, name, trait_spec, value):
         if name == 'out_file':
             return ''
-        else:
-            return super(WarpPoints, self)._format_arg(name, trait_spec, value)
+
+        return super(WarpPoints, self)._format_arg(name, trait_spec, value)
 
     def _parse_inputs(self, skip=None):
         fname, ext = op.splitext(self.inputs.in_coords)
@@ -1894,21 +1896,12 @@ def _parse_inputs(self, skip=None):
 
     def _vtk_to_coords(self, in_file, out_file=None):
         # Ensure that tvtk is loaded with the appropriate ETS_TOOLKIT env var
-        old_ets = os.getenv('ETS_TOOLKIT')
-        os.environ['ETS_TOOLKIT'] = 'null'
-        try:
-            from tvtk.api import tvtk
-            from tvtk.tvtk_classes.vtk_version import vtk_build_version
-        except ImportError:
-            vtk_build_version = None
-            raise ImportError('This interface requires tvtk to run.')
-        finally:
-            if old_ets is not None:
-                os.environ['ETS_TOOLKIT'] = old_ets
-            else:
-                del os.environ['ETS_TOOLKIT']
+        if VTKInfo.no_tvtk():
+            raise ImportError('TVTK is required and tvtk package was not found')
 
-        vtk_major = int(vtk_build_version[0])
+        from ...algorithms.mesh import tvtk
+
+        vtk_major = VTKInfo.vtk_version()[0]
         reader = tvtk.PolyDataReader(file_name=in_file + '.vtk')
         reader.update()
 
@@ -1923,23 +1916,13 @@ def _vtk_to_coords(self, in_file, out_file=None):
 
     def _coords_to_vtk(self, points, out_file):
         import os
-
         # Ensure that tvtk is loaded with the appropriate ETS_TOOLKIT env var
-        old_ets = os.getenv('ETS_TOOLKIT')
-        os.environ['ETS_TOOLKIT'] = 'null'
-        try:
-            from tvtk.api import tvtk
-            from tvtk.tvtk_classes.vtk_version import vtk_build_version
-        except ImportError:
-            vtk_build_version = None
-            raise ImportError('This interface requires tvtk to run.')
-        finally:
-            if old_ets is not None:
-                os.environ['ETS_TOOLKIT'] = old_ets
-            else:
-                del os.environ['ETS_TOOLKIT']
+        if VTKInfo.no_tvtk():
+            raise ImportError('TVTK is required and tvtk package was not found')
+
+        from ...algorithms.mesh import tvtk
 
-        vtk_major = int(vtk_build_version[0])
+        vtk_major = VTKInfo.vtk_version()[0]
         reader = tvtk.PolyDataReader(file_name=self.inputs.in_file)
         reader.update()