From 9c8fcdcf6fbde7afb611cecf0682378595893b4e Mon Sep 17 00:00:00 2001
From: git-abhishek <36498066+git-abhishek@users.noreply.github.com>
Date: Wed, 27 Jun 2018 17:12:43 -0400
Subject: [PATCH] Update as per first code review

---
 yt/visualization/image_writer.py                 | 13 +++++++------
 yt/visualization/tests/test_image_writer.py      | 13 +++++++++++--
 yt/visualization/tests/test_offaxisprojection.py | 14 +++++++++-----
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/yt/visualization/image_writer.py b/yt/visualization/image_writer.py
index f72bb73953..606efb21b8 100644
--- a/yt/visualization/image_writer.py
+++ b/yt/visualization/image_writer.py
@@ -22,7 +22,7 @@
     get_image_suffix, \
     get_brewer_cmap
 from yt.units.yt_array import YTQuantity
-from yt.utilities.exceptions import YTNotInsideNotebook, YTException
+from yt.utilities.exceptions import YTNotInsideNotebook
 from .color_maps import mcm
 from . import _colormap_data as cmd
 import yt.utilities.lib.image_utilities as au
@@ -150,9 +150,8 @@ def write_bitmap(bitmap_array, filename, max_val = None, transpose=False):
         first element will be placed in the lower-left corner.
     """
     if len(bitmap_array.shape) != 3 or bitmap_array.shape[-1] not in (3, 4):
-        raise YTException(message="Expecting image array of shape (N,M,3)"
-                                  " or (N,M,4), received %s"
-                                  % str(bitmap_array.shape))
+        raise RuntimeError("Expecting image array of shape (N,M,3) or "
+                           "(N,M,4), received %s" % str(bitmap_array.shape))
 
     if bitmap_array.dtype != np.uint8:
         s1, s2 = bitmap_array.shape[:2]
@@ -298,8 +297,10 @@ def strip_colormap_data(fn = "color_map_data.py",
     f.write("### Auto-generated colormap tables, taken from Matplotlib ###\n\n")
     f.write("from numpy import array\n")
     f.write("color_map_luts = {}\n\n\n")
-    if cmaps is None: cmaps = rcm.ColorMaps
-    if isinstance(cmaps, string_types): cmaps = [cmaps]
+    if cmaps is None:
+        cmaps = rcm.ColorMaps
+    if isinstance(cmaps, string_types):
+        cmaps = [cmaps]
     for cmap_name in sorted(cmaps):
         vals = rcm._extract_lookup_table(cmap_name)
         f.write("### %s ###\n\n" % (cmap_name))
diff --git a/yt/visualization/tests/test_image_writer.py b/yt/visualization/tests/test_image_writer.py
index 93f7b68a4f..3b0b967407 100644
--- a/yt/visualization/tests/test_image_writer.py
+++ b/yt/visualization/tests/test_image_writer.py
@@ -1,3 +1,13 @@
+"""
+Tests for visualization.image_writer
+"""
+# -----------------------------------------------------------------------------
+# Copyright (c) 2018, yt Development Team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file COPYING.txt, distributed with this software.
+# -----------------------------------------------------------------------------
 import os
 import shutil
 import tempfile
@@ -7,7 +17,6 @@
 from nose.tools import assert_raises
 
 from yt.testing import fake_random_ds, assert_equal
-from yt.utilities.exceptions import YTException
 from yt.visualization.image_writer import multi_image_composite, splat_points, \
     write_bitmap, apply_colormap, strip_colormap_data
 
@@ -55,7 +64,7 @@ def test_write_bitmap(self):
         png_str_trans = write_bitmap(image_trans, None, transpose=True)
         assert_equal(png_str, png_str_trans)
 
-        with assert_raises(YTException) as ex:
+        with assert_raises(RuntimeError) as ex:
             write_bitmap(np.ones([16, 16]), None)
         desired = ("Expecting image array of shape (N,M,3) "
                    "or (N,M,4), received (16, 16)")
diff --git a/yt/visualization/tests/test_offaxisprojection.py b/yt/visualization/tests/test_offaxisprojection.py
index 5b08473cab..659c478727 100644
--- a/yt/visualization/tests/test_offaxisprojection.py
+++ b/yt/visualization/tests/test_offaxisprojection.py
@@ -11,16 +11,15 @@
 # Distributed under the terms of the Modified BSD License.
 #
 # The full license is in the file COPYING.txt, distributed with this software.
-#-----------------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 import os
-import os.path
 import shutil
 import tempfile
 import unittest
 
-from yt.mods import write_projection
+from yt.visualization.image_writer import write_projection
 from yt.testing import \
-    fake_random_ds, assert_equal, expand_keywords
+    fake_random_ds, expand_keywords, assert_fname
 from yt.visualization.volume_rendering.api import off_axis_projection
 
 
@@ -78,9 +77,14 @@ def test_oap(self):
             image = off_axis_projection(*oap_args, **oap_kwargs)
             for wp_kwargs in wp_kwargs_list:
                 write_projection(image, fn % i, **wp_kwargs)
-                assert_equal(os.path.exists(fn % i), True)
+                assert_fname(fn % i)
 
         # Test remaining parameters of write_projection
         write_projection(image, "test_2", xlabel="x-axis", ylabel="y-axis")
+        assert_fname("test_2.png")
+
         write_projection(image, "test_3.pdf", xlabel="x-axis", ylabel="y-axis")
+        assert_fname("test_3.pdf")
+
         write_projection(image, "test_4.eps", xlabel="x-axis", ylabel="y-axis")
+        assert_fname("test_4.eps")