Merge branch 'soql-flag'

cmd/root.go

@@ -18,64 +18,102 @@ func init() {
 	RootCmd.Flags().BoolP("write", "w", false, "write result to (source) file instead of stdout")
 	RootCmd.Flags().BoolP("list", "l", false, "list files whose formatting differs from apexfmt's")
 	RootCmd.Flags().BoolP("verbose", "v", false, "enable debug logging")
+	RootCmd.Flags().BoolP("soql", "s", false, "format SOQL query")
+
+	RootCmd.MarkFlagsMutuallyExclusive("write", "list")
+	RootCmd.MarkFlagsMutuallyExclusive("soql", "write")
+	RootCmd.MarkFlagsMutuallyExclusive("soql", "list")
+
 }
 
 var RootCmd = &cobra.Command{
 	Use:   "apexfmt [file...]",
 	Short: "Format Apex",
-	Run: func(cmd *cobra.Command, args []string) {
+	RunE: func(cmd *cobra.Command, args []string) error {
+		if soql, _ := cmd.Flags().GetBool("soql"); soql {
+			formatSOQL()
+			return nil
+		}
+
 		write, _ := cmd.Flags().GetBool("write")
 		list, _ := cmd.Flags().GetBool("list")
 		verbose, _ := cmd.Flags().GetBool("verbose")
 		if verbose {
 			log.SetLevel(log.DebugLevel)
 		}
+		formatters := []*formatter.Formatter{}
 		for _, filename := range args {
-			f := formatter.NewFormatter(filename)
+			formatters = append(formatters, formatter.NewFormatter(filename, nil))
+		}
+		if len(args) == 0 {
+			if write {
+				return fmt.Errorf("One or more files required for --write option")
+			}
+			if list {
+				return fmt.Errorf("One or more files required for --list option")
+			}
+			formatters = append(formatters, formatter.NewFormatter("", os.Stdin))
+		}
+		for _, f := range formatters {
 			err := f.Format()
 			if err != nil {
-				fmt.Fprintf(os.Stderr, "Failed to format file %s: %s\n", filename, err.Error())
+				fmt.Fprintf(os.Stderr, "Failed to format file %s: %s\n", f.SourceName(), err.Error())
 				os.Exit(1)
 			}
 
 			if list {
 				changed, err := f.Changed()
 				if err != nil {
-					fmt.Fprintf(os.Stderr, "Failed to check file for changes %s: %s\n", filename, err.Error())
+					fmt.Fprintf(os.Stderr, "Failed to check file for changes %s: %s\n", f.SourceName(), err.Error())
 					os.Exit(1)
 				}
 				if changed {
-					fmt.Println(filename)
+					fmt.Println(f.SourceName())
 				}
 			} else if !write {
 				out, err := f.Formatted()
 				if err != nil {
-					fmt.Fprintf(os.Stderr, "Failed to get formatted source %s: %s\n", filename, err.Error())
+					fmt.Fprintf(os.Stderr, "Failed to get formatted source %s: %s\n", f.SourceName(), err.Error())
 					os.Exit(1)
 				}
 				fmt.Println(out)
 			}
 			changed, err := f.Changed()
 			if err != nil {
-				fmt.Fprintf(os.Stderr, "Failed to check file for changes %s: %s\n", filename, err.Error())
+				fmt.Fprintf(os.Stderr, "Failed to check file for changes %s: %s\n", f.SourceName(), err.Error())
 				os.Exit(1)
 			}
 			if write && changed {
 				err = f.Write()
 				if err != nil {
-					fmt.Fprintf(os.Stderr, "Failed to write file %s: %s\n", filename, err.Error())
+					fmt.Fprintf(os.Stderr, "Failed to write file %s: %s\n", f.SourceName(), err.Error())
 					os.Exit(1)
 				}
 			}
 		}
+		return nil
 	},
-	Args:                  cobra.MinimumNArgs(1),
 	DisableFlagsInUseLine: true,
 }
 
 func globalConfig() {
 }
 
+func formatSOQL() {
+	f := formatter.NewSOQLFormatter()
+	err := f.Format()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to format query: %s\n", err.Error())
+		os.Exit(1)
+	}
+	out, err := f.Formatted()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to get formatted query: %s\n", err.Error())
+		os.Exit(1)
+	}
+	fmt.Println(out)
+}
+
 func Execute() {
 	if err := RootCmd.Execute(); err != nil {
 		fmt.Println(err)

docs/apexfmt.md

@@ -9,8 +9,10 @@ apexfmt [file...]
 ### Options
 
 ```
-  -h, --help    help for apexfmt
-  -l, --list    list files whose formatting differs from apexfmt's
-  -w, --write   write result to (source) file instead of stdout
+  -h, --help      help for apexfmt
+  -l, --list      list files whose formatting differs from apexfmt's
+  -s, --soql      format SOQL query
+  -v, --verbose   enable debug logging
+  -w, --write     write result to (source) file instead of stdout
 ```
 

formatter/formatter.go

@@ -13,23 +13,40 @@ import (
 
 type Formatter struct {
 	filename  string
+	reader    io.Reader
 	source    []byte
 	formatted []byte
 }
 
+func (f *Formatter) SourceName() string {
+	if f.filename != "" {
+		return f.filename
+	}
+	return "<stdin>"
+}
+
 type errorListener struct {
 	*antlr.DefaultErrorListener
 	filename string
 }
 
 func (e *errorListener) SyntaxError(_ antlr.Recognizer, _ interface{}, line, column int, msg string, _ antlr.RecognitionException) {
-	_, _ = fmt.Fprintln(os.Stderr, e.filename+" line "+strconv.Itoa(line)+":"+strconv.Itoa(column)+" "+msg)
+	if e.filename == "" {
+		_, _ = fmt.Fprintln(os.Stderr, "line "+strconv.Itoa(line)+":"+strconv.Itoa(column)+" "+msg)
+	} else {
+		_, _ = fmt.Fprintln(os.Stderr, e.filename+" line "+strconv.Itoa(line)+":"+strconv.Itoa(column)+" "+msg)
+	}
 	os.Exit(1)
 }
 
-func NewFormatter(filename string) *Formatter {
+func NewFormatter(filename string, reader io.Reader) *Formatter {
+	if filename != "" {
+		return &Formatter{
+			filename: filename,
+		}
+	}
 	return &Formatter{
-		filename: filename,
+		reader: reader,
 	}
 }
 
@@ -55,9 +72,9 @@ func (f *Formatter) Changed() (bool, error) {
 
 func (f *Formatter) Format() error {
 	if f.source == nil {
-		src, err := readFile(f.filename)
+		src, err := readFile(f.filename, f.reader)
 		if err != nil {
-			return fmt.Errorf("Failed to read file %s: %w", f.filename, err)
+			return fmt.Errorf("Failed to read file %s: %w", f.SourceName(), err)
 		}
 		f.source = src
 	}
@@ -86,13 +103,17 @@ func (f *Formatter) Write() error {
 	return writeFile(f.filename, f.formatted)
 }
 
-func readFile(filename string) ([]byte, error) {
-	f, err := os.Open(filename)
-	if err != nil {
-		return nil, err
+func readFile(filename string, reader io.Reader) ([]byte, error) {
+	r := reader
+	if filename != "" {
+		f, err := os.Open(filename)
+		if err != nil {
+			return nil, err
+		}
+		defer f.Close()
+		r = f
 	}
-	defer f.Close()
-	src, err := io.ReadAll(f)
+	src, err := io.ReadAll(r)
 	if err != nil {
 		return nil, err
 	}

formatter/soql.go

@@ -0,0 +1,65 @@
+package formatter
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/antlr4-go/antlr/v4"
+	"github.com/octoberswimmer/apexfmt/parser"
+)
+
+type SOQLFormatter struct {
+	source    []byte
+	formatted []byte
+}
+
+func NewSOQLFormatter() *SOQLFormatter {
+	return &SOQLFormatter{}
+}
+
+func (f *SOQLFormatter) Formatted() (string, error) {
+	if f.formatted == nil {
+		err := f.Format()
+		if err != nil {
+			return "", err
+		}
+	}
+	return string(f.formatted), nil
+}
+
+func (f *SOQLFormatter) Changed() (bool, error) {
+	if f.formatted == nil {
+		err := f.Format()
+		if err != nil {
+			return false, err
+		}
+	}
+	return !bytes.Equal(f.source, f.formatted), nil
+}
+
+func (f *SOQLFormatter) Format() error {
+	if f.source == nil {
+		src, err := io.ReadAll(os.Stdin)
+		if err != nil {
+			return fmt.Errorf("Failed to read in query: %w", err)
+		}
+		f.source = src
+	}
+	input := antlr.NewInputStream(string(f.source))
+	lexer := parser.NewApexLexer(input)
+	stream := antlr.NewCommonTokenStream(lexer, antlr.TokenDefaultChannel)
+
+	p := parser.NewApexParser(stream)
+	p.RemoveErrorListeners()
+	p.AddErrorListener(&errorListener{})
+
+	v := NewFormatVisitor(stream)
+	out, ok := v.visitRule(p.Query()).(string)
+	if !ok {
+		return fmt.Errorf("Unexpected result parsing apex")
+	}
+	f.formatted = append([]byte(out), '\n')
+	return nil
+}

formatter/visitors.go

@@ -674,6 +674,11 @@ func (v *FormatVisitor) VisitSoqlLiteral(ctx *parser.SoqlLiteralContext) interfa
 }
 
 func (v *FormatVisitor) VisitQuery(ctx *parser.QueryContext) interface{} {
+	i := NewChainVisitor()
+	n := i.visitRule(ctx).(int)
+	if n > 3 {
+		defer restoreWrap(wrap(v))
+	}
 	sep := " "
 	indent := 0
 	if v.wrap {