Functions now support optional arguments.

Author: springcomp

jmespath/functions.py

@@ -81,22 +81,28 @@ def call_function(self, function_name, resolved_args):
         return function(self, *resolved_args)
 
     def _validate_arguments(self, args, signature, function_name):
-        if signature and signature[-1].get('variadic'):
+        required_arguments_count = len([param for param in signature if not param.get('optional') or not param['optional']])
+        optional_arguments_count = len([param for param in signature if param.get('optional') and param['optional']])
+        has_variadic = signature[-1].get('variadic') if signature != None else False
+        if has_variadic:
             if len(args) < len(signature):
                 raise exceptions.VariadictArityError(
                     len(signature), len(args), function_name)
-        elif len(args) != len(signature):
+        elif optional_arguments_count > 0:
+            if len(args) < required_arguments_count or len(args) > (required_arguments_count + optional_arguments_count):
+                raise exceptions.ArityError(
+                    len(signature), len(args), function_name)
+        elif len(args) != required_arguments_count:
             raise exceptions.ArityError(
                 len(signature), len(args), function_name)
         return self._type_check(args, signature, function_name)
 
     def _type_check(self, actual, signature, function_name):
-        for i in range(len(signature)):
-            allowed_types = signature[i]['types']
+        for i in range(min(len(signature), len(actual))):
+            allowed_types = self._get_allowed_types_from_signature(signature[i])
             if allowed_types:
                 self._type_check_single(actual[i], allowed_types,
                                         function_name)
-
     def _type_check_single(self, current, types, function_name):
         # Type checking involves checking the top level type,
         # and in the case of arrays, potentially checking the types
@@ -120,6 +126,13 @@ def _type_check_single(self, current, types, function_name):
             self._subtype_check(current, allowed_subtypes,
                                 types, function_name)
 
+    ## signature supports monotype {'type': 'type-name'}
+    ## or multiple types {'types': ['type1-name', 'type2-name']}
+    def _get_allowed_types_from_signature(self, spec):
+        if spec.get('type'):
+            spec.update({'types': [spec.get('type')]})
+        return spec.get('types')
+
     def _get_allowed_pytypes(self, types):
         allowed_types = []
         allowed_subtypes = []

tests/test_signatures.py

@@ -0,0 +1,69 @@
+from tests import unittest
+
+from jmespath import compat
+from jmespath import exceptions
+from jmespath import functions
+
+class TestFunctionSignatures(unittest.TestCase):
+
+	def setUp(self):
+		self._functions = functions.Functions()
+
+	def test_signature_with_monotype_argument(self):
+		(function_name, signature) = self._make_test("_function_with_monotyped_arguments")
+		self._functions._validate_arguments(['string'], signature, function_name)
+		self.assertRaises(
+			exceptions.ArityError, lambda :
+			self._functions._validate_arguments([], signature, function_name)
+		)
+
+	def test_signature_with_optional_arguments(self):
+		(function_name, signature) = self._make_test("_function_with_optional_arguments")
+		self._functions._validate_arguments(['string'], signature, function_name)
+		self._functions._validate_arguments(['string', 42], signature, function_name)
+		self._functions._validate_arguments(['string', 43], signature, function_name)
+		self.assertRaises(
+			exceptions.ArityError, lambda :
+			self._functions._validate_arguments([], signature, function_name)
+		)
+		self.assertRaises(
+			exceptions.ArityError, lambda :
+			self._functions._validate_arguments(['string', 42, 43, 44], signature, function_name)
+		)
+
+	def test_signature_with_variadic_arguments(self):
+		(function_name, signature) = self._make_test("_function_with_variadic_arguments")
+		self._functions._validate_arguments(['string', 'text1'], signature, function_name)
+		self._functions._validate_arguments(['string', 'text1', 'text2'], signature, function_name)
+		self.assertRaises(
+			exceptions.VariadictArityError, lambda :
+			self._functions._validate_arguments(['string'], signature, function_name)
+		)
+
+	def _make_test(self, funcName):
+		for name, method in compat.get_methods(TestFunctionSignatures):
+			print(name)
+			if name != funcName:
+				continue
+			signature = getattr(method, 'signature', None)
+			return (funcName, signature)
+		return None
+
+	## arg1 can only be of type 'string'
+	## this signature supports testing simplified syntax
+	## where 'type' is a string instead of an array of strings
+	@functions.signature({'type': 'string'})
+	def _function_with_monotyped_arguments(self, arg1):
+		return None
+
+	@functions.signature({'type': 'string'}, {'type': 'string', 'variadic': True})
+	def _function_with_variadic_arguments(self, arg1, *arguments):
+		return None
+
+	@functions.signature({'type': 'string'}, {'type': 'number', 'optional': True}, {'type': 'number', 'optional': True})
+	def _function_with_optional_arguments(self, arg1, opt1, opt2):
+		return None
+
+
+if __name__ == '__main__':
+    unittest.main()