diff --git a/jmespath/functions.py b/jmespath/functions.py index 627b569d..21769bbc 100644 --- a/jmespath/functions.py +++ b/jmespath/functions.py @@ -235,10 +235,16 @@ def _func_reverse(self, arg): @signature({"types": ['number']}) def _func_ceil(self, arg): + if math.isnan(arg) or math.isinf(arg): + raise exceptions.JMESPathTypeError( + 'ceil', arg, 'non-finite number', ['number']) return math.ceil(arg) @signature({"types": ['number']}) def _func_floor(self, arg): + if math.isnan(arg) or math.isinf(arg): + raise exceptions.JMESPathTypeError( + 'floor', arg, 'non-finite number', ['number']) return math.floor(arg) @signature({"types": ['string']}, {"types": ['array-string']}) diff --git a/tests/test_functions.py b/tests/test_functions.py index a352d69e..f9ec4305 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -55,3 +55,25 @@ def test_variadic_is_pluralized(self): self.assertEqual( str(exception), 'Expected at least 1 argument for function not_null(), received 0') + + def test_ceil_raises_jmespath_error_on_infinity(self): + # Regression test: ceil() used to pass its argument straight to + # math.ceil(), which raises the builtin OverflowError on infinity + # instead of a JMESPathError subclass. + with self.assertRaises(exceptions.JMESPathTypeError): + jmespath.search('ceil(@)', float('inf')) + + def test_ceil_raises_jmespath_error_on_nan(self): + with self.assertRaises(exceptions.JMESPathTypeError): + jmespath.search('ceil(@)', float('nan')) + + def test_floor_raises_jmespath_error_on_infinity(self): + # Regression test: floor() used to pass its argument straight to + # math.floor(), which raises the builtin OverflowError on infinity + # instead of a JMESPathError subclass. + with self.assertRaises(exceptions.JMESPathTypeError): + jmespath.search('floor(@)', float('inf')) + + def test_floor_raises_jmespath_error_on_nan(self): + with self.assertRaises(exceptions.JMESPathTypeError): + jmespath.search('floor(@)', float('nan'))