fix: restore support for JFunction and add test case

rayokota · Sep 16, 2024 · fa0a5fd · fa0a5fd
1 parent 13cc9c7
commit fa0a5fd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/jsonata/functions.py b/src/jsonata/functions.py
@@ -1575,7 +1575,7 @@ def get_function_arity(func: Any) -> int:
 
             return len(signature(func.function).parameters)
         else:
-            raise RuntimeError(f"unexpected function '{type(func)}'")
+            return len(func.arguments)
 
     #
     # Helper function to build the arguments to be supplied to the function arg of the

diff --git a/tests/custom_function_test.py b/tests/custom_function_test.py
@@ -23,7 +23,11 @@ def test_ternary(self):
         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(self):
+    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]