fix: support keep array workaround for map

Author: netomi

src/jsonata/parser.py

@@ -1180,6 +1180,7 @@ def lambda1(terms):
                 result.position = expr.position
                 result.lhs = self.process_ast(expr.lhs)
                 result.rhs = self.process_ast(expr.rhs)
+                result.keep_array = result.lhs.keep_array or result.rhs.keep_array
             else:
                 result = Parser.Infix(self, None)
                 result.type = expr.type
@@ -1245,6 +1246,7 @@ def lambda4(arg):
 
             result.arguments = [lambda4(x) for x in expr.arguments]
             result.procedure = self.process_ast(expr.procedure)
+            result.keep_array = result.procedure.keep_array
         elif type == "lambda":
             result = Parser.Symbol(self)
             result.type = expr.type

tests/custom_function_test.py

@@ -31,3 +31,13 @@ def test_map_with_lambda(self):
     def test_map_with_function(self):
         expression = jsonata.Jsonata("$map([1, 2, 3], function($v) { $v * $v })")
         assert expression.evaluate(None) == [1, 4, 9]
+
+    def test_singleton_map(self):
+        expression = jsonata.Jsonata("$map([1], $square)[]")
+        expression.register_lambda("square", lambda x: x * x)
+        assert expression.evaluate(None) == [1]
+
+    def test_singleton_map_with_chain(self):
+        expression = jsonata.Jsonata("$ ~> $map($square)[]")
+        expression.register_lambda("square", lambda x: x * x)
+        assert expression.evaluate([1]) == [1]