From da5e2b99492e2f9ffed8d61507656b0f6860ab1b Mon Sep 17 00:00:00 2001
From: Marco Tezzele <marcotez@gmail.com>
Date: Tue, 12 Apr 2016 17:49:21 +0200
Subject: [PATCH] vtk plot

---
 pygem/vtkhandler.py      | 49 ++++++++++++++++++++++++++++++++++++++++
 tests/test_vtkhandler.py |  7 ++++++
 2 files changed, 56 insertions(+)

diff --git a/pygem/vtkhandler.py b/pygem/vtkhandler.py
index db47c067..894c5327 100644
--- a/pygem/vtkhandler.py
+++ b/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()
+
+		
\ No newline at end of file
diff --git a/tests/test_vtkhandler.py b/tests/test_vtkhandler.py
index 95df6aa1..657028be 100644
--- a/tests/test_vtkhandler.py
+++ b/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)
+