From 86bb90f9890b66a05b89425bd1fbda0392d4e1cf Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Fri, 16 Feb 2024 10:47:50 +1100
Subject: [PATCH] Can shebang yq scripts! #1851

---
 acceptance_tests/shebang.sh | 26 ++++++++++++++++++++++++++
 cmd/utils.go                | 11 ++++++++++-
 examples/environment.yq     |  1 +
 test.yq                     |  2 ++
 4 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100755 acceptance_tests/shebang.sh
 mode change 100644 => 100755 examples/environment.yq
 create mode 100755 test.yq

diff --git a/acceptance_tests/shebang.sh b/acceptance_tests/shebang.sh
new file mode 100755
index 00000000..5d47e751
--- /dev/null
+++ b/acceptance_tests/shebang.sh
@@ -0,0 +1,26 @@
+#!/bin/bash
+
+setUp() {
+  rm test*.yq || true
+  cat >test.yq <<EOL
+#!./yq
+.a.b
+EOL
+chmod +x test.yq
+
+rm test*.yml || true
+  cat >test.yml <<EOL
+a: {b: apple}
+EOL
+}
+
+testCanExecYqFile() {
+  read -r -d '' expected << EOM
+apple
+EOM
+   X=$(./test.yq test.yml)
+  assertEquals "$expected" "$X"
+}
+
+source ./scripts/shunit2
+
diff --git a/cmd/utils.go b/cmd/utils.go
index 16ddcf9e..b266016d 100644
--- a/cmd/utils.go
+++ b/cmd/utils.go
@@ -255,6 +255,16 @@ func processStdInArgs(args []string) []string {
 
 func processArgs(originalArgs []string) (string, []string, error) {
 	expression := forceExpression
+	args := processStdInArgs(originalArgs)
+	maybeFirstArgIsAFile := len(args) > 0 && maybeFile(args[0])
+
+	if expressionFile == "" && maybeFirstArgIsAFile && strings.HasSuffix(args[0], ".yq") {
+		// lets check if an expression file was given
+		yqlib.GetLogger().Debug("Assuming arg %v is an expression file", args[0])
+		expressionFile = args[0]
+		args = args[1:]
+	}
+
 	if expressionFile != "" {
 		expressionBytes, err := os.ReadFile(expressionFile)
 		if err != nil {
@@ -264,7 +274,6 @@ func processArgs(originalArgs []string) (string, []string, error) {
 		expression = strings.ReplaceAll(string(expressionBytes), "\r\n", "\n")
 	}
 
-	args := processStdInArgs(originalArgs)
 	yqlib.GetLogger().Debugf("processed args: %v", args)
 	if expression == "" && len(args) > 0 && args[0] != "-" && !maybeFile(args[0]) {
 		yqlib.GetLogger().Debug("assuming expression is '%v'", args[0])
diff --git a/examples/environment.yq b/examples/environment.yq
old mode 100644
new mode 100755
index 764f5802..81325364
--- a/examples/environment.yq
+++ b/examples/environment.yq
@@ -1,3 +1,4 @@
+#!yq
 .[] |(
     ( select(kind == "scalar") | key + "='" + . + "'"),
     ( select(kind == "seq") | key + "=(" + (map("'" + . + "'") | join(",")) + ")")
diff --git a/test.yq b/test.yq
new file mode 100755
index 00000000..0b1150b9
--- /dev/null
+++ b/test.yq
@@ -0,0 +1,2 @@
+#!./yq
+.a.b