From 17ff39ff251dc8bd14540ada99506836a1d02554 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Sat, 20 Sep 2014 17:12:35 -0500 Subject: [PATCH] Allow assert_nD to take multiple dimensions --- skimage/_shared/utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/skimage/_shared/utils.py b/skimage/_shared/utils.py index da64a32d..805f3cce 100644 --- a/skimage/_shared/utils.py +++ b/skimage/_shared/utils.py @@ -145,6 +145,9 @@ def safe_as_int(val, atol=1e-3): def assert_nD(array, arg_name='image', ndim=2): array = np.asanyarray(array) - if array.ndim != ndim: + msg = "The parameter `%s` must be a %s-dimensional array" + if isinstance(ndim, int): + ndim = [ndim] + if not array.ndim in ndim: msg = "The parameter `%s` must be a %s-dimensional array" - raise ValueError(msg % (arg_name, ndim)) + raise ValueError(msg % (arg_name, '-or-'.join([str(n) for n in ndim]))