Skip to content

Commit 5b63c49

Browse files
committed
_check_ax_type function
1 parent fd0af86 commit 5b63c49

File tree

2 files changed

+19
-3
lines changed

2 files changed

+19
-3
lines changed

probscale/tests/test_probscale/test_viz.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,17 @@ def test_linlog(self):
270270
)
271271

272272

273+
class Test__check_ax_type(object):
274+
@nt.raises(ValueError)
275+
def test_bad_value(self):
276+
viz._check_ax_type("JUNK")
277+
278+
def test_upper(self):
279+
nt.assert_equal('pp', viz._check_ax_type('PP'))
280+
nt.assert_equal('qq', viz._check_ax_type('QQ'))
281+
nt.assert_equal('prob', viz._check_ax_type('ProB'))
282+
283+
273284
@image_comparison(baseline_images=['test_probplot_prob'], extensions=['png'])
274285
def test_probplot_prob():
275286
fig, ax = plt.subplots()

probscale/viz.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,13 @@ def _check_ax_name(axname, argname):
4040
msg = 'Invalid value for {} ({}). Must be on of {}.'
4141
raise ValueError(msg.format(argname, arg, valid_args))
4242

43-
return axname
43+
return axname.lower()
44+
45+
46+
def _check_ax_type(axtype):
47+
if axtype.lower() not in ['pp', 'qq', 'prob']:
48+
raise ValueError("invalid axtype: {}".format(axtype))
49+
return axtype.lower()
4450

4551

4652
def probplot(data, ax=None, axtype='prob', probax='x',
@@ -98,8 +104,7 @@ def probplot(data, ax=None, axtype='prob', probax='x',
98104
line_kws = {} if line_kws is None else line_kws.copy()
99105

100106
# check axtype
101-
if axtype not in ['pp', 'qq', 'prob']:
102-
raise ValueError("invalid axtype: {}".format(axtype))
107+
axtype = _check_ax_type(axtype)
103108

104109
# compute the plotting positions and sort the data
105110
qntls, datavals = stats.probplot(data, fit=False)

0 commit comments

Comments
 (0)