diff --git a/jmespath/functions.py b/jmespath/functions.py index 627b569d..f6866909 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 not math.isfinite(arg): + raise exceptions.JMESPathTypeError( + 'ceil', arg, 'non-finite number', ['number']) return math.ceil(arg) @signature({"types": ['number']}) def _func_floor(self, arg): + if not math.isfinite(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..74f38db5 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -55,3 +55,24 @@ def test_variadic_is_pluralized(self): self.assertEqual( str(exception), 'Expected at least 1 argument for function not_null(), received 0') + + def test_floor_ceil_reject_non_finite(self): + # json.loads accepts Infinity/NaN by default, so data loaded the + # usual way can reach floor()/ceil(). math.floor/ceil raise + # uncatchable OverflowError/ValueError on non-finite floats; we + # surface them as JMESPathTypeError so callers catching + # JMESPathError can handle them like any other bad argument. + for expression, document in [ + ('floor(@)', 'Infinity'), + ('ceil(@)', 'Infinity'), + ('floor(@)', 'NaN'), + ('ceil(@)', 'NaN'), + ]: + with self.assertRaises(exceptions.JMESPathTypeError): + jmespath.search(expression, json.loads(document)) + + def test_floor_ceil_still_work_for_finite(self): + self.assertEqual(jmespath.search('floor(@)', 3.7), 3) + self.assertEqual(jmespath.search('ceil(@)', 3.2), 4) + self.assertEqual(jmespath.search('floor(`3.7`)', {}), 3) + self.assertEqual(jmespath.search('ceil(`3.2`)', {}), 4)