From d9af2dd97641f8cccdef65f759123cafd437fc38 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Sat, 15 Jun 2024 21:56:32 +1000
Subject: [PATCH] Support env[] expression #2071

---
 pkg/yqlib/lexer_participle.go   |  2 +-
 pkg/yqlib/operation.go          |  4 ++--
 pkg/yqlib/operator_env_test.go  | 10 ++++++++++
 pkg/yqlib/operator_keys_test.go | 10 ++++++++++
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/pkg/yqlib/lexer_participle.go b/pkg/yqlib/lexer_participle.go
index ef31e778f5c..af288999a43 100644
--- a/pkg/yqlib/lexer_participle.go
+++ b/pkg/yqlib/lexer_participle.go
@@ -406,7 +406,7 @@ func envOp(strenv bool) yqAction {
 		envOperation.OperationType = envOpType
 		envOperation.Preferences = preferences
 
-		return &token{TokenType: operationToken, Operation: envOperation}, nil
+		return &token{TokenType: operationToken, Operation: envOperation, CheckForPostTraverse: envOpType.CheckForPostTraverse}, nil
 	}
 }
 
diff --git a/pkg/yqlib/operation.go b/pkg/yqlib/operation.go
index d79ac7a0944..3dcabd89476 100644
--- a/pkg/yqlib/operation.go
+++ b/pkg/yqlib/operation.go
@@ -161,7 +161,7 @@ var stringInterpolationOpType = &operationType{Type: "STRING_INT", NumArgs: 0, P
 var loadOpType = &operationType{Type: "LOAD", NumArgs: 1, Precedence: 52, Handler: loadOperator}
 var loadStringOpType = &operationType{Type: "LOAD_STRING", NumArgs: 1, Precedence: 52, Handler: loadStringOperator}
 
-var keysOpType = &operationType{Type: "KEYS", NumArgs: 0, Precedence: 50, Handler: keysOperator}
+var keysOpType = &operationType{Type: "KEYS", NumArgs: 0, Precedence: 52, Handler: keysOperator, CheckForPostTraverse: true}
 
 var collectObjectOpType = &operationType{Type: "COLLECT_OBJECT", NumArgs: 0, Precedence: 50, Handler: collectObjectOperator}
 
@@ -175,7 +175,7 @@ var traverseArrayOpType = &operationType{Type: "TRAVERSE_ARRAY", NumArgs: 2, Pre
 var selfReferenceOpType = &operationType{Type: "SELF", NumArgs: 0, Precedence: 55, Handler: selfOperator}
 var valueOpType = &operationType{Type: "VALUE", NumArgs: 0, Precedence: 50, Handler: valueOperator, ToString: valueToStringFunc}
 var referenceOpType = &operationType{Type: "REF", NumArgs: 0, Precedence: 50, Handler: referenceOperator}
-var envOpType = &operationType{Type: "ENV", NumArgs: 0, Precedence: 50, Handler: envOperator}
+var envOpType = &operationType{Type: "ENV", NumArgs: 0, Precedence: 52, Handler: envOperator, CheckForPostTraverse: true}
 var notOpType = &operationType{Type: "NOT", NumArgs: 0, Precedence: 50, Handler: notOperator}
 var toNumberOpType = &operationType{Type: "TO_NUMBER", NumArgs: 0, Precedence: 50, Handler: toNumberOperator}
 var emptyOpType = &operationType{Type: "EMPTY", Precedence: 50, Handler: emptyOperator}
diff --git a/pkg/yqlib/operator_env_test.go b/pkg/yqlib/operator_env_test.go
index 878c5bcec97..ea7e1c6046b 100644
--- a/pkg/yqlib/operator_env_test.go
+++ b/pkg/yqlib/operator_env_test.go
@@ -5,6 +5,16 @@ import (
 )
 
 var envOperatorScenarios = []expressionScenario{
+	{
+		description:          "Read string environment variable",
+		skipDoc:              true,
+		environmentVariables: map[string]string{"myenv": "[cat,dog]"},
+		expression:           `env(myenv)[]`,
+		expected: []string{
+			"D0, P[0], (!!str)::cat\n",
+			"D0, P[1], (!!str)::dog\n",
+		},
+	},
 	{
 		description:          "Read string environment variable",
 		environmentVariables: map[string]string{"myenv": "cat meow"},
diff --git a/pkg/yqlib/operator_keys_test.go b/pkg/yqlib/operator_keys_test.go
index b0936608413..15086c8e23a 100644
--- a/pkg/yqlib/operator_keys_test.go
+++ b/pkg/yqlib/operator_keys_test.go
@@ -39,6 +39,16 @@ var keysOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!seq)::- dog\n- cat\n",
 		},
 	},
+	{
+		description: "Map keys with splat",
+		skipDoc:     true,
+		document:    `{dog: woof, cat: meow}`,
+		expression:  `keys[]`,
+		expected: []string{
+			"D0, P[dog], (!!str)::dog\n",
+			"D0, P[cat], (!!str)::cat\n",
+		},
+	},
 	{
 		skipDoc:    true,
 		document:   `{}`,