Merge pull request #6 from netomi/fix-arity-of-lambda-functions

fix: correctly calculate arity of lambda functions

Author: rayokota

src/jsonata/functions.py

@@ -1570,8 +1570,11 @@ def get_function_arity(func: Any) -> int:
         from jsonata import jsonata
         if isinstance(func, jsonata.Jsonata.JFunction):
             return func.signature.get_min_number_of_args()
+        elif isinstance(func, jsonata.Jsonata.JLambda):
+            from inspect import signature
+
+            return len(signature(func.function).parameters)
         else:
-            # Lambda
             return len(func.arguments)
 
     #

tests/custom_function_test.py

@@ -22,3 +22,12 @@ def test_ternary(self):
         expression = jsonata.Jsonata("$abc(a,b,c)")
         expression.register_lambda("abc", lambda x, y, z: str(x) + str(y) + str(z))
         assert expression.evaluate({"a": "a", "b": "b", "c": "c"}) == "abc"
+
+    def test_map_with_lambda(self):
+        expression = jsonata.Jsonata("$map([1, 2, 3], $square)")
+        expression.register_lambda("square", lambda x: x * x)
+        assert expression.evaluate(None) == [1, 4, 9]
+
+    def test_map_with_function(self):
+        expression = jsonata.Jsonata("$map([1, 2, 3], function($v) { $v * $v })")
+        assert expression.evaluate(None) == [1, 4, 9]