autogenpb/protoReformat.go

package main

import (
	"fmt"
	"os"
	"strings"

	"go.wit.com/log"
)

// like 'goimport' but for .proto files

var maxVarname int
var maxVartype int

func protoReformat(filename string) error {
	// read in the .proto file
	data, err := os.ReadFile(filename)
	if err != nil {
		log.Info("file read failed", filename, err)
		return err
	}

	pf, err := os.OpenFile(filename+".new", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Info("file open error. permissions?", filename, err)
		return err
	}
	defer pf.Close()

	var inMessage bool
	var curmsg []string

	// gets the max vartype and varname
	for _, line := range strings.Split(string(data), "\n") {
		if strings.HasPrefix(line, "message ") {
			inMessage = true
			continue
		}

		// find the end of the message
		if strings.HasPrefix(line, "}") {
			inMessage = false
			formatMessage(curmsg)
			curmsg = nil
			continue
		}

		// don't format or change anything when not in a "message {" section
		if !inMessage {
			continue
		}
		curmsg = append(curmsg, line)
	}

	// parse the proto file for message struct names
	for _, line := range strings.Split(string(data), "\n") {
		if strings.HasPrefix(line, "message ") {
			inMessage = true
			parts := strings.Fields(line)
			if len(parts) > 3 {
				start := parts[0] + " " + parts[1] + " " + parts[2]
				end := strings.Join(parts[3:], " ")
				offset := maxVarname + maxVartype + 16 - len(start)
				pad := fmt.Sprintf("%d", offset)
				hmm := "%s     %" + pad + "s %s"
				line = fmt.Sprintf(hmm, start, " ", end)
			}
			fmt.Fprintln(pf, line)
			continue
		}

		// find the end of the message
		if strings.HasPrefix(line, "}") {
			inMessage = false
			for _, newline := range formatMessage(curmsg) {
				fmt.Fprintln(pf, newline)
			}
			fmt.Fprintln(pf, line)
			curmsg = nil
			continue
		}

		// don't format or change anything when not in a "message {" section
		if !inMessage {
			fmt.Fprintln(pf, line)
			continue
		}
		curmsg = append(curmsg, line)
	}

	// for i, s := range slices.Backward(pf.ToSort) {
	return nil
}

func formatMessage(curmsg []string) []string {
	var newmsg []string

	// find the max length of varname and vartype
	for _, line := range curmsg {
		parts := strings.Split(line, ";")
		if len(parts) < 2 {
			log.Info("parse error", line)
			return curmsg
		}

		vartype, varname, _, _ := tokenMsgVar(line)
		if len(vartype) > maxVartype {
			maxVartype = len(vartype)
		}
		if len(varname) > maxVarname {
			maxVarname = len(varname)
		}
	}

	for _, line := range curmsg {
		mt := fmt.Sprintf("%d", maxVartype)
		mv := fmt.Sprintf("%d", maxVarname)

		hmm := "        %-" + mt + "s   %-" + mv + "s     = %-3s %s"

		vartype, varname, id, end := tokenMsgVar(line)
		end = strings.TrimSpace(end)
		id = id + ";"

		newline := fmt.Sprintf(hmm, vartype, varname, id, end)
		newmsg = append(newmsg, newline)
	}
	return newmsg
}

// returns vartype, varname, id, end
func tokenMsgVar(line string) (string, string, string, string) {
	parts := strings.Split(line, ";")
	front := parts[0]
	end := strings.Join(parts[1:], ";")

	var id string
	var varname string
	var vartype string

	parts = strings.Fields(front)
	parts, id = slicesPop(parts)
	parts, _ = slicesPop(parts) // this is the "=" sign
	parts, varname = slicesPop(parts)
	vartype = strings.Join(parts, " ")

	return vartype, varname, id, end
}

func slicesPop(parts []string) ([]string, string) {
	if len(parts) == 0 {
		return nil, ""
	}
	if len(parts) == 1 {
		return nil, parts[0]
	}
	x := len(parts)
	end := parts[x-1]
	return parts[0 : x-1], end
}