From efae1938fd6c8434532ca7527cd90752e558d377 Mon Sep 17 00:00:00 2001
From: duxinlong <duxinlong@bytedance.com>
Date: Wed, 8 Feb 2023 12:01:48 +0000
Subject: [PATCH] feat: support more env than terminal

Change-Id: I7f35e90b8f19f4ea781832885d35e2f1e275207a
---
 parse.go | 31 +++++++++++++++++++++++++++----
 usage.go |  6 +++---
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/parse.go b/parse.go
index b935fad..6d8b509 100644
--- a/parse.go
+++ b/parse.go
@@ -5,6 +5,7 @@ import (
 	"encoding/csv"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"reflect"
@@ -119,6 +120,10 @@ type Config struct {
 	// StrictSubcommands intructs the library not to allow global commands after
 	// subcommand
 	StrictSubcommands bool
+
+	OsExit func(int)
+	Stdout io.Writer
+	Stderr io.Writer
 }
 
 // Parser represents a set of command line options with destination values
@@ -132,6 +137,10 @@ type Parser struct {
 
 	// the following field changes during processing of command line arguments
 	lastCmd *command
+
+	osExit func(int)
+	stdout io.Writer
+	stderr io.Writer
 }
 
 // Versioned is the interface that the destination struct should implement to
@@ -196,6 +205,20 @@ func NewParser(config Config, dests ...interface{}) (*Parser, error) {
 	p := Parser{
 		cmd:    &command{name: name},
 		config: config,
+
+		osExit: osExit,
+		stdout: stdout,
+		stderr: stderr,
+	}
+
+	if config.OsExit != nil {
+		p.osExit = config.OsExit
+	}
+	if config.Stdout != nil {
+		p.stdout = config.Stdout
+	}
+	if config.Stderr != nil {
+		p.stderr = config.Stderr
 	}
 
 	// make a list of roots
@@ -483,11 +506,11 @@ func (p *Parser) MustParse(args []string) {
 	err := p.Parse(args)
 	switch {
 	case err == ErrHelp:
-		p.writeHelpForSubcommand(stdout, p.lastCmd)
-		osExit(0)
+		p.writeHelpForSubcommand(p.stdout, p.lastCmd)
+		p.osExit(0)
 	case err == ErrVersion:
-		fmt.Fprintln(stdout, p.version)
-		osExit(0)
+		fmt.Fprintln(p.stdout, p.version)
+		p.osExit(0)
 	case err != nil:
 		p.failWithSubcommand(err.Error(), p.lastCmd)
 	}
diff --git a/usage.go b/usage.go
index 7a480c3..80eba45 100644
--- a/usage.go
+++ b/usage.go
@@ -39,9 +39,9 @@ func (p *Parser) FailSubcommand(msg string, subcommand ...string) error {
 
 // failWithSubcommand prints usage information for the given subcommand to stderr and exits with non-zero status
 func (p *Parser) failWithSubcommand(msg string, cmd *command) {
-	p.writeUsageForSubcommand(stderr, cmd)
-	fmt.Fprintln(stderr, "error:", msg)
-	osExit(-1)
+	p.writeUsageForSubcommand(p.stderr, cmd)
+	fmt.Fprintln(p.stderr, "error:", msg)
+	p.osExit(-1)
 }
 
 // WriteUsage writes usage information to the given writer