recover panic and move state to its own func

pressly · Apr 10, 2024 · 836ab6f · 836ab6f
1 parent 5e65c72
commit 836ab6f
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 37 deletions.
diff --git a/internal/cli/cli.go b/internal/cli/cli.go
@@ -3,7 +3,6 @@ package cli
 import (
 	"context"
 	"fmt"
-	"io/fs"
 	"os"
 	"os/signal"
 	"runtime"
@@ -17,17 +16,15 @@ import (
 // return.
 func Main(opts ...Options) {
 	ctx, stop := newContext()
-	defer stop()
-	defer func() {
-		if r := recover(); r != nil {
-			fmt.Fprintf(os.Stderr, "unexpected error: %v\n", r)
+	go func() {
+		defer stop()
+		if err := Run(ctx, os.Args[1:], opts...); err != nil {
+			fmt.Fprintln(os.Stderr, err.Error())
 			os.Exit(1)
 		}
 	}()
-	if err := Run(ctx, os.Args[1:], opts...); err != nil {
-		fmt.Fprintln(os.Stderr, err.Error())
-		os.Exit(1)
-	}
+	// TODO(mf): this looks wonky because we're not waiting for the context to be done. But
+	// eventually, I'd like to add a timeout here so we don't hang indefinitely.
 	<-ctx.Done()
 	os.Exit(0)
 }
@@ -38,30 +35,7 @@ func Main(opts ...Options) {
 // Options can be used to customize the behavior of the CLI, such as setting the environment,
 // redirecting stdout and stderr, and providing a custom filesystem such as embed.FS.
 func Run(ctx context.Context, args []string, opts ...Options) error {
-	state := &state{
-		environ: os.Environ(),
-	}
-	for _, opt := range opts {
-		if err := opt.apply(state); err != nil {
-			return err
-		}
-	}
-	// Set defaults if not set by the caller.
-	if state.stdout == nil {
-		state.stdout = os.Stdout
-	}
-	if state.stderr == nil {
-		state.stderr = os.Stderr
-	}
-	if state.fsys == nil {
-		// Use the default filesystem if not set, reading from the local filesystem.
-		state.fsys = func(dir string) (fs.FS, error) { return os.DirFS(dir), nil }
-	}
-	if state.openConnection == nil {
-		// Use the default openConnection function if not set.
-		state.openConnection = openConnection
-	}
-	return run(ctx, state, args)
+	return run(ctx, args, opts...)
 }
 
 func newContext() (context.Context, context.CancelFunc) {

diff --git a/internal/cli/run.go b/internal/cli/run.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io/fs"
+	"os"
 	"strings"
 
 	"github.com/peterbourgon/ff/v4"
@@ -13,7 +15,17 @@ const (
 	ENV_NO_COLOR = "NO_COLOR"
 )
 
-func run(ctx context.Context, st *state, args []string) error {
+func run(ctx context.Context, args []string, opts ...Options) (retErr error) {
+	defer func() {
+		if r := recover(); r != nil {
+			retErr = fmt.Errorf("panic: %v", r)
+		}
+	}()
+	st, err := newStateWithDefaults(opts...)
+	if err != nil {
+		return err
+	}
+
 	root := newRootCommand(st)
 	// Add subcommands
 	commands := []func(*state) (*ff.Command, error){
@@ -28,11 +40,10 @@ func run(ctx context.Context, st *state, args []string) error {
 	}
 
 	// Parse the flags and return help if requested.
-	err := root.Parse(
+	if err := root.Parse(
 		args,
 		ff.WithEnvVarPrefix("GOOSE"), // Support environment variables for all flags
-	)
-	if err != nil {
+	); err != nil {
 		if errors.Is(err, ff.ErrHelp) {
 			fmt.Fprintf(st.stderr, "\n%s\n", createHelp(root))
 			return nil
@@ -47,6 +58,33 @@ func run(ctx context.Context, st *state, args []string) error {
 	return root.Run(ctx)
 }
 
+func newStateWithDefaults(opts ...Options) (*state, error) {
+	state := &state{
+		environ: os.Environ(),
+	}
+	for _, opt := range opts {
+		if err := opt.apply(state); err != nil {
+			return nil, err
+		}
+	}
+	// Set defaults if not set by the caller
+	if state.stdout == nil {
+		state.stdout = os.Stdout
+	}
+	if state.stderr == nil {
+		state.stderr = os.Stderr
+	}
+	if state.fsys == nil {
+		// Use the default filesystem if not set, reading from the local filesystem.
+		state.fsys = func(dir string) (fs.FS, error) { return os.DirFS(dir), nil }
+	}
+	if state.openConnection == nil {
+		// Use the default openConnection function if not set.
+		state.openConnection = openConnection
+	}
+	return state, nil
+}
+
 func checkRequiredFlags(cmd *ff.Command) error {
 	if cmd != nil {
 		cmd = cmd.GetSelected()