From e976e985efa0a192fd8f8f771301a320388cb244 Mon Sep 17 00:00:00 2001
From: Mitsuo Heijo <mitsuo.heijo@gmail.com>
Date: Fri, 4 Dec 2020 23:33:01 +0900
Subject: [PATCH] Add a new function to parse a line containing environment
 variables

---
 README.md          |  6 ++++++
 shellwords.go      | 29 +++++++++++++++++++++++++++++
 shellwords_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+)

diff --git a/README.md b/README.md
index e91902f..8840672 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,12 @@ args, err := shellwords.Parse("./foo --bar=baz")
 // args should be ["./foo", "--bar=baz"]
 ```
 
+```go
+envs, args, err := shellwords.ParseWithEnvs("FOO=foo BAR=baz ./foo --bar=baz")
+// envs should be ["FOO=foo", "BAR=baz"]
+// args should be ["./foo", "--bar=baz"]
+```
+
 ```go
 os.Setenv("FOO", "bar")
 p := shellwords.NewParser()
diff --git a/shellwords.go b/shellwords.go
index fc29fa4..4f11960 100644
--- a/shellwords.go
+++ b/shellwords.go
@@ -290,6 +290,35 @@ loop:
 	return args, nil
 }
 
+func (p *Parser) ParseWithEnvs(line string) (envs []string, args []string, err error) {
+	_args, err := p.Parse(line)
+	if err != nil {
+		return nil, nil, err
+	}
+	envs = []string{}
+	args = []string{}
+	parsingEnv := true
+	for _, arg := range _args {
+		if parsingEnv && isEnv(arg) {
+			envs = append(envs, arg)
+		} else {
+			if parsingEnv {
+				parsingEnv = false
+			}
+			args = append(args, arg)
+		}
+	}
+	return envs, args, nil
+}
+
+func isEnv(arg string) bool {
+	return len(strings.Split(arg, "=")) == 2
+}
+
 func Parse(line string) ([]string, error) {
 	return NewParser().Parse(line)
 }
+
+func ParseWithEnvs(line string) (envs []string, args []string, err error) {
+	return NewParser().ParseWithEnvs(line)
+}
diff --git a/shellwords_test.go b/shellwords_test.go
index b32a493..4f94528 100644
--- a/shellwords_test.go
+++ b/shellwords_test.go
@@ -407,3 +407,45 @@ func TestEnvInQuoted(t *testing.T) {
 		t.Fatalf("Expected %#v, but %#v:", expected, args)
 	}
 }
+
+func TestParseWithEnvs(t *testing.T) {
+	tests := []struct {
+		line               string
+		wantEnvs, wantArgs []string
+	}{
+		{
+			line:     "FOO=foo cmd --args=A=B",
+			wantEnvs: []string{"FOO=foo"},
+			wantArgs: []string{"cmd", "--args=A=B"},
+		},
+		{
+			line:     "FOO=foo BAR=bar cmd --args=A=B -A=B",
+			wantEnvs: []string{"FOO=foo", "BAR=bar"},
+			wantArgs: []string{"cmd", "--args=A=B", "-A=B"},
+		},
+		{
+			line:     `sh -c "FOO=foo BAR=bar cmd --args=A=B -A=B"`,
+			wantEnvs: []string{},
+			wantArgs: []string{"sh", "-c", "FOO=foo BAR=bar cmd --args=A=B -A=B"},
+		},
+		{
+			line:     "cmd --args=A=B -A=B",
+			wantEnvs: []string{},
+			wantArgs: []string{"cmd", "--args=A=B", "-A=B"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.line, func(t *testing.T) {
+			envs, args, err := ParseWithEnvs(tt.line)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !reflect.DeepEqual(envs, tt.wantEnvs) {
+				t.Errorf("Expected %#v, but %#v", tt.wantEnvs, envs)
+			}
+			if !reflect.DeepEqual(args, tt.wantArgs) {
+				t.Errorf("Expected %#v, but %#v", tt.wantArgs, args)
+			}
+		})
+	}
+}