From 14875b34e31faada4f7672000abf91bc49ce0de1 Mon Sep 17 00:00:00 2001
From: Alon Krymgand <me@alon.kr>
Date: Mon, 12 Aug 2024 11:30:33 +0300
Subject: [PATCH] tidy main

---
 usm.go | 112 ++++++++++++++++++++++++++++++++++-----------------------
 1 file changed, 67 insertions(+), 45 deletions(-)

diff --git a/usm.go b/usm.go
index f68e634..c4ea532 100644
--- a/usm.go
+++ b/usm.go
@@ -4,65 +4,87 @@ import (
 	"fmt"
 	"os"
 
-	"github.com/spf13/cobra"
-
 	"alon.kr/x/usm/lex"
 	"alon.kr/x/usm/parse"
 	"alon.kr/x/usm/source"
+	"github.com/spf13/cobra"
 )
 
+func setInputSource(cmd *cobra.Command, args []string) error {
+	if len(args) > 0 {
+		file, err := os.Open(args[0])
+		if err != nil {
+			return fmt.Errorf("error opening file: %v", err)
+		}
+		cmd.SetIn(file)
+	}
+	return nil
+}
+
+func lexCommand(cmd *cobra.Command, args []string) {
+	view, err := source.ReadSource(cmd.InOrStdin())
+	if err != nil {
+		fmt.Printf("Error reading source: %v\n", err)
+		os.Exit(1)
+	}
+
+	tokens, err := lex.NewTokenizer().Tokenize(view)
+	if err != nil {
+		fmt.Printf("Error tokenizing: %v\n", err)
+		os.Exit(1)
+	}
+
+	_, ctx := view.Detach()
+	for _, tkn := range tokens {
+		fmt.Printf("%s ", tkn.String(ctx))
+		if tkn.Type == lex.SeparatorToken {
+			fmt.Println()
+		}
+	}
+}
+
+func fmtCommand(cmd *cobra.Command, args []string) {
+	view, err := source.ReadSource(cmd.InOrStdin())
+	if err != nil {
+		fmt.Printf("Error reading source: %v\n", err)
+		os.Exit(1)
+	}
+
+	tokens, err := lex.NewTokenizer().Tokenize(view)
+	if err != nil {
+		fmt.Printf("Error tokenizing: %v\n", err)
+		os.Exit(1)
+	}
+
+	tknView := parse.NewTokenView(tokens)
+	file, perr := parse.NewFileParser().Parse(&tknView)
+	if perr == nil {
+		fmt.Print(file.String(view.Ctx()))
+	} else {
+		fmt.Println(perr.Error(view.Ctx()))
+	}
+}
+
 func main() {
 	rootCmd := &cobra.Command{
 		Use:   "usm",
-		Short: " One Universal assembly language to rule them all.",
+		Short: "One Universal assembly language to rule them all.",
 	}
 
 	lexCmd := &cobra.Command{
-		Use:   "lex",
-		Short: "Lex the source code",
-		Run: func(cmd *cobra.Command, args []string) {
-			view, err := source.ReadSource(os.Stdin)
-			if err != nil {
-				panic(err)
-			}
-
-			tokens, err := lex.NewTokenizer().Tokenize(view)
-			if err != nil {
-				panic(err)
-			}
-
-			_, ctx := view.Detach()
-			for _, tkn := range tokens {
-				fmt.Printf("%s ", tkn.String(ctx))
-				if tkn.Type == lex.SeparatorToken {
-					fmt.Println()
-				}
-			}
-		},
+		Use:     "lex [file]",
+		Short:   "Lex the source code",
+		Args:    cobra.MaximumNArgs(1),
+		PreRunE: setInputSource,
+		Run:     lexCommand,
 	}
 
 	fmtCmd := &cobra.Command{
-		Use:   "fmt",
-		Short: "Format the source code",
-		Run: func(cmd *cobra.Command, args []string) {
-			view, err := source.ReadSource(os.Stdin)
-			if err != nil {
-				panic(err)
-			}
-
-			tokens, err := lex.NewTokenizer().Tokenize(view)
-			if err != nil {
-				panic(err)
-			}
-
-			tknView := parse.NewTokenView(tokens)
-			file, perr := parse.NewFileParser().Parse(&tknView)
-			if perr == nil {
-				fmt.Print(file.String(view.Ctx()))
-			} else {
-				fmt.Println(perr.Error(view.Ctx()))
-			}
-		},
+		Use:     "fmt [file]",
+		Short:   "Format the source code",
+		Args:    cobra.MaximumNArgs(1),
+		PreRunE: setInputSource,
+		Run:     fmtCommand,
 	}
 
 	rootCmd.AddCommand(lexCmd)