From 1d54e2cd8f73cd89cde04a5108b928e2c8e9d14a Mon Sep 17 00:00:00 2001
From: Chetan Sarva <chetan@pixelcop.net>
Date: Fri, 6 Dec 2024 14:38:07 -0500
Subject: [PATCH] feat: allow running command from util

---
 go.mod  |  2 +-
 main.go | 75 +++++++++++++++++++++++++++++++++++++++++++++++----------
 2 files changed, 63 insertions(+), 14 deletions(-)

diff --git a/go.mod b/go.mod
index d221980..ffa26a3 100644
--- a/go.mod
+++ b/go.mod
@@ -1,5 +1,5 @@
 module github.com/chetan/elapsed
 
-go 1.17
+go 1.23
 
 require github.com/mitchellh/go-homedir v1.1.0
diff --git a/main.go b/main.go
index 730ac7e..1436866 100644
--- a/main.go
+++ b/main.go
@@ -4,9 +4,12 @@ import (
 	"bufio"
 	"flag"
 	"fmt"
+	"io"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/mitchellh/go-homedir"
@@ -77,18 +80,7 @@ func parseFlags() {
 	}
 }
 
-func main() {
-	parseFlags()
-	if !(showElapsed || showDelta) {
-		fmt.Println("error: must enable either elapsed or delta")
-		os.Exit(1)
-	}
-	stat, _ := os.Stdin.Stat()
-	if (stat.Mode() & os.ModeCharDevice) != 0 {
-		fmt.Println("error: stdin not connected to a pipe")
-		os.Exit(1)
-	}
-
+func run(r io.Reader) {
 	f := "["
 	if showTS != "" {
 		f += "%s"
@@ -107,7 +99,7 @@ func main() {
 	}
 	f += "] %s"
 
-	reader := bufio.NewReader(os.Stdin)
+	reader := bufio.NewReader(r)
 	start := time.Now()
 	last := time.Now()
 	for {
@@ -151,3 +143,60 @@ func main() {
 		}
 	}
 }
+
+func getReader() (io.Reader, []string) {
+	stat, _ := os.Stdin.Stat()
+	if (stat.Mode() & os.ModeCharDevice) != 0 {
+		// look for cmd
+		args := os.Args[1:]
+		if len(args) != 0 {
+			for i, arg := range args {
+				if arg == "--" && i+1 < len(args) && args[i+1] != "" {
+					return nil, args[i+1:]
+				}
+			}
+		}
+
+		fmt.Println("error: stdin not connected to a pipe")
+		os.Exit(1)
+	}
+
+	return os.Stdin, nil
+}
+
+func runCmd(args []string) {
+	cmd := exec.Command(args[0], args[1:]...)
+	r, w := io.Pipe()
+	cmd.Stdout = w
+	cmd.Stderr = w
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		err := cmd.Start()
+		if err != nil {
+			fmt.Println("error running command: ", err)
+		}
+		go run(r)
+		err = cmd.Wait()
+		if err != nil {
+			fmt.Println("error waiting for command: ", err)
+		}
+		os.Exit(cmd.ProcessState.ExitCode())
+	}()
+	wg.Wait()
+}
+
+func main() {
+	parseFlags()
+	if !(showElapsed || showDelta) {
+		fmt.Println("error: must enable either elapsed or delta")
+		os.Exit(1)
+	}
+
+	r, args := getReader()
+	if args != nil {
+		runCmd(args)
+	} else {
+		run(r)
+	}
+}