Raise JMESPathTypeError for non-finite floor()/ceil() args

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']})

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)