vtk plot

pygem/vtkhandler.py

@@ -93,3 +93,52 @@ def write(self, mesh_points, filename):
 			writer.SetInputData(data)
 
 		writer.Write()
+
+
+	def plot(self, plot_file=None):
+		"""
+		Method to plot an stl file. If `plot_file` is not given it plots `self.infile`.
+
+		:param string plot_file: the stl filename you want to plot.
+		"""
+		if plot_file is None:
+			plot_file = self.infile
+		else:
+			self._check_filename_type(plot_file)
+
+		# Read the source file.
+		reader = vtk.vtkUnstructuredGridReader()
+		reader.SetFileName(plot_file)
+		reader.Update() # Needed because of GetScalarRange
+		output = reader.GetOutput()
+		scalar_range = output.GetScalarRange()
+
+		# Create the mapper that corresponds the objects of the vtk file
+		# into graphics elements
+		mapper = vtk.vtkDataSetMapper()
+		if vtk.VTK_MAJOR_VERSION <= 5:
+			mapper.SetInput(output)
+		else:
+			mapper.SetInputData(output)
+		mapper.SetScalarRange(scalar_range)
+
+		# Create the Actor
+		actor = vtk.vtkActor()
+		actor.SetMapper(mapper)
+
+		# Create the Renderer
+		renderer = vtk.vtkRenderer()
+		renderer.AddActor(actor)
+		renderer.SetBackground(20, 20, 20) # Set background color (white is 1, 1, 1)
+
+		# Create the RendererWindow
+		renderer_window = vtk.vtkRenderWindow()
+		renderer_window.AddRenderer(renderer)
+
+		# Create the RendererWindowInteractor and display the vtk_file
+		interactor = vtk.vtkRenderWindowInteractor()
+		interactor.SetRenderWindow(renderer_window)
+		interactor.Initialize()
+		interactor.Start()
+
+

tests/test_vtkhandler.py

@@ -133,3 +133,10 @@ def test_vtk_write_comparison(self):
 		vtk_handler.write(mesh_points, outfilename)
 		self.assertTrue(filecmp.cmp(outfilename, outfilename_expected))
 		os.remove(outfilename)
+
+
+	def test_vtk_plot_failing_outfile_type(self):
+		vtk_handler = vh.VtkHandler()
+		with self.assertRaises(TypeError):
+			vtk_handler.plot(plot_file=1.1)
+