parametrized tests and cleaner calls to validate

Vedant Nanda · Vedant Nanda · commit 04ebb0d82b40 · 2017-04-04T23:12:17.000+05:30
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -585,6 +585,19 @@ def _determine_lims(self, xmin=None, xmax=None, *args, **kwargs):
             xmax += 0.05
         return (xmin, xmax)
 
+    def _validate_axis_limits(self, limit, convert):
+        """
+        If the axis limits being set are infinite, this function
+
+        raises an error.
+
+        """
+        if limit is not None:
+            if (isinstance(limit, float) and
+                    (not np.isreal(limit) or not np.isfinite(limit))):
+                raise ValueError("NaN or Inf cannot be the argument values")
+            return convert(limit)
+
     def set_xlim3d(self, left=None, right=None, emit=True, auto=False, **kw):
         """
         Set 3D x limits.
@@ -603,10 +616,8 @@ def set_xlim3d(self, left=None, right=None, emit=True, auto=False, **kw):
             left, right = left
 
         self._process_unit_info(xdata=(left, right))
-        if left is not None:
-            left = self.convert_xunits(left)
-        if right is not None:
-            right = self.convert_xunits(right)
+        left = self._validate_axis_limits(left, self.convert_xunits)
+        right = self._validate_axis_limits(right, self.convert_xunits)
 
         old_left, old_right = self.get_xlim()
         if left is None:
@@ -663,10 +674,8 @@ def set_ylim3d(self, bottom=None, top=None, emit=True, auto=False, **kw):
             top = self.convert_yunits(top)
 
         old_bottom, old_top = self.get_ylim()
-        if bottom is None:
-            bottom = old_bottom
-        if top is None:
-            top = old_top
+        bottom = self._validate_axis_limits(bottom, self.convert_yunits)
+        top = self._validate_axis_limits(top, self.convert_yunits)
 
         if top == bottom:
             warnings.warn(('Attempting to set identical bottom==top results\n'
@@ -711,10 +720,8 @@ def set_zlim3d(self, bottom=None, top=None, emit=True, auto=False, **kw):
             bottom, top = bottom
 
         self._process_unit_info(zdata=(bottom, top))
-        if bottom is not None:
-            bottom = self.convert_zunits(bottom)
-        if top is not None:
-            top = self.convert_zunits(top)
+        bottom = self._validate_axis_limits(bottom, self.convert_yunits)
+        top = self._validate_axis_limits(top, self.convert_yunits)
 
         old_bottom, old_top = self.get_zlim()
         if bottom is None:
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -513,3 +513,20 @@ def test_autoscale():
     ax.set_autoscalez_on(True)
     ax.plot([0, 2], [0, 2], [0, 2])
     assert ax.get_w_lims() == (0, 1, -.1, 1.1, -.4, 2.4)
+
+
+@pytest.mark.parametrize('value', [np.inf, np.nan])
+@pytest.mark.parametrize(('setter', 'side'), [
+    ('set_xlim3d', 'left'),
+    ('set_xlim3d', 'right'),
+    ('set_ylim3d', 'bottom'),
+    ('set_ylim3d', 'top'),
+    ('set_zlim3d', 'bottom'),
+    ('set_zlim3d', 'top'),
+])
+def test_invalid_axes_limits(setter, side, value):
+    limit = {side: value}
+    with pytest.raises(ValueError):
+        fig = plt.figure()
+        obj = fig.add_subplot(111, projection='3d')
+        getattr(obj, setter)(**limit)
