From d3c758468a90663b0ffdf9da0ff9ab94c255259c Mon Sep 17 00:00:00 2001 From: Thomas Neidhart Date: Mon, 16 Sep 2024 23:27:14 +0200 Subject: [PATCH] fix: support keep array workaround for map --- src/jsonata/parser.py | 2 ++ tests/custom_function_test.py | 10 ++++++++++ 2 files changed, 12 insertions(+) diff --git a/src/jsonata/parser.py b/src/jsonata/parser.py index 8ccaf88..186dae1 100644 --- a/src/jsonata/parser.py +++ b/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 diff --git a/tests/custom_function_test.py b/tests/custom_function_test.py index ef2426a..d481b62 100644 --- a/tests/custom_function_test.py +++ b/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]