use the validation module in viz.py

Author: phobson

probscale/tests/test_probscale/test_viz.py

@@ -24,61 +24,6 @@ def setup_plot_data():
     return data
 
 
-class Test__check_ax_obj(object):
-    @nt.raises(ValueError)
-    def test_bad_value(self):
-        viz._check_ax_obj('junk')
-
-    @cleanup
-    def test_with_ax(self):
-        fig, ax = plt.subplots()
-        fig1, ax1 = viz._check_ax_obj(ax)
-        nt.assert_true(isinstance(ax1, plt.Axes))
-        nt.assert_true(isinstance(fig1, plt.Figure))
-        nt.assert_true(ax1 is ax)
-        nt.assert_true(fig1 is fig)
-
-    @cleanup
-    def test_with_None(self):
-        fig1, ax1 = viz._check_ax_obj(None)
-        nt.assert_true(isinstance(ax1, plt.Axes))
-        nt.assert_true(isinstance(fig1, plt.Figure))
-
-
-class Test__check_fit_arg(object):
-    @nt.raises(ValueError)
-    def test_bad_fitarg(self):
-        viz._check_fit_arg('junk', 'fitprobs')
-
-    def test_x(self):
-        nt.assert_equal('x', viz._check_fit_arg('x', 'fitprobs'))
-        nt.assert_equal('x', viz._check_fit_arg('x', 'fitlogs'))
-
-    def test_y(self):
-        nt.assert_equal('y', viz._check_fit_arg('y', 'fitprobs'))
-        nt.assert_equal('y', viz._check_fit_arg('y', 'fitlogs'))
-
-    def test_both(self):
-        nt.assert_equal('both', viz._check_fit_arg('both', 'fitprobs'))
-        nt.assert_equal('both', viz._check_fit_arg('both', 'fitlogs'))
-
-    def test_None(self):
-        nt.assert_true(viz._check_fit_arg(None, 'fitprobs') is None)
-        nt.assert_true(viz._check_fit_arg(None, 'fitlogs') is None)
-
-
-class Test__check_ax_name(object):
-    @nt.raises
-    def test_bad_name(self):
-        viz._check_fit_arg('junk', 'axname')
-
-    def test_x(self):
-        nt.assert_equal('x', viz._check_fit_arg('x', 'axname'))
-
-    def test_y(self):
-        nt.assert_equal('y', viz._check_fit_arg('y', 'axname'))
-
-
 class Test__fit_line(object):
     def setup(self):
         self.data = np.array([
@@ -271,17 +216,6 @@ def test_linlog(self):
         )
 
 
-class Test__check_ax_type(object):
-    @nt.raises(ValueError)
-    def test_bad_value(self):
-        viz._check_ax_type("JUNK")
-
-    def test_upper(self):
-        nt.assert_equal('pp', viz._check_ax_type('PP'))
-        nt.assert_equal('qq', viz._check_ax_type('QQ'))
-        nt.assert_equal('prob', viz._check_ax_type('ProB'))
-
-
 @image_comparison(baseline_images=['test_probplot_prob'], extensions=['png'])
 def test_probplot_prob():
     fig, ax = plt.subplots()

probscale/viz.py

@@ -3,52 +3,14 @@
 from matplotlib import scale
 from scipy import stats
 
-from .probscale import ProbScale, _minimal_norm
+from .probscale import ProbScale
+from .probscale import _minimal_norm
+from . import validate
 
 
 scale.register_scale(ProbScale)
 
 
-def _check_ax_obj(ax):
-    """ Checks if a value if an Axes. If None, a new one is created.
-
-    """
-
-    if ax is None:
-        fig, ax = pyplot.subplots()
-    elif isinstance(ax, pyplot.Axes):
-        fig = ax.figure
-    else:
-        msg = "`ax` must be a matplotlib Axes instance or None"
-        raise ValueError(msg)
-
-    return fig, ax
-
-
-def _check_fit_arg(arg, argname):
-    valid_args = ['x', 'y', 'both', None]
-    if arg not in valid_args:
-        msg = 'Invalid value for {} ({}). Must be on of {}.'
-        raise ValueError(msg.format(argname, arg, valid_args))
-
-    return arg
-
-
-def _check_ax_name(axname, argname):
-    valid_args = ['x', 'y']
-    if axname.lower() not in valid_args:
-        msg = 'Invalid value for {} ({}). Must be on of {}.'
-        raise ValueError(msg.format(argname, arg, valid_args))
-
-    return axname.lower()
-
-
-def _check_ax_type(axtype):
-    if axtype.lower() not in ['pp', 'qq', 'prob']:
-        raise ValueError("invalid axtype: {}".format(axtype))
-    return axtype.lower()
-
-
 def probplot(data, ax=None, axtype='prob', probax='x',
              otherscale='linear', xlabel=None, ylabel=None,
              bestfit=False, return_results=False,
@@ -99,15 +61,15 @@ def probplot(data, ax=None, axtype='prob', probax='x',
     """
 
     # check input values
-    fig, ax = _check_ax_obj(ax)
-    probax = _check_ax_name(probax, 'probax')
+    fig, ax = validate.axes_object(ax)
+    probax = validate.axis_name(probax, 'x')
 
     # default values for plotting options
     scatter_kws = {} if scatter_kws is None else scatter_kws.copy()
     line_kws = {} if line_kws is None else line_kws.copy()
 
     # check axtype
-    axtype = _check_ax_type(axtype)
+    axtype = validate.axis_type(axtype)
 
     # compute the plotting positions and sort the data
     qntls, datavals = stats.probplot(data, fit=False)
@@ -205,8 +167,8 @@ def _fit_line(x, y, xhat=None, fitprobs=None, fitlogs=None, dist=None):
 
     """
 
-    fitprobs = _check_fit_arg(fitprobs, "fitprobs")
-    fitlogs = _check_fit_arg(fitlogs, "fitlogs")
+    fitprobs = validate.fit_argument(fitprobs, "fitprobs")
+    fitlogs = validate.fit_argument(fitlogs, "fitlogs")
 
     # maybe set xhat to default values
     if xhat is None: