printer.go raw

   1  package printer
   2  
   3  import (
   4  	"fmt"
   5  	"math"
   6  	"strings"
   7  
   8  	"github.com/fatih/color"
   9  	"github.com/goccy/go-yaml/ast"
  10  	"github.com/goccy/go-yaml/token"
  11  )
  12  
  13  // Property additional property set for each the token
  14  type Property struct {
  15  	Prefix string
  16  	Suffix string
  17  }
  18  
  19  // PrintFunc returns property instance
  20  type PrintFunc func() *Property
  21  
  22  // Printer create text from token collection or ast
  23  type Printer struct {
  24  	LineNumber       bool
  25  	LineNumberFormat func(num int) string
  26  	MapKey           PrintFunc
  27  	Anchor           PrintFunc
  28  	Alias            PrintFunc
  29  	Bool             PrintFunc
  30  	String           PrintFunc
  31  	Number           PrintFunc
  32  }
  33  
  34  func defaultLineNumberFormat(num int) string {
  35  	return fmt.Sprintf("%2d | ", num)
  36  }
  37  
  38  func (p *Printer) property(tk *token.Token) *Property {
  39  	prop := &Property{}
  40  	switch tk.PreviousType() {
  41  	case token.AnchorType:
  42  		if p.Anchor != nil {
  43  			return p.Anchor()
  44  		}
  45  		return prop
  46  	case token.AliasType:
  47  		if p.Alias != nil {
  48  			return p.Alias()
  49  		}
  50  		return prop
  51  	}
  52  	switch tk.NextType() {
  53  	case token.MappingValueType:
  54  		if p.MapKey != nil {
  55  			return p.MapKey()
  56  		}
  57  		return prop
  58  	}
  59  	switch tk.Type {
  60  	case token.BoolType:
  61  		if p.Bool != nil {
  62  			return p.Bool()
  63  		}
  64  		return prop
  65  	case token.AnchorType:
  66  		if p.Anchor != nil {
  67  			return p.Anchor()
  68  		}
  69  		return prop
  70  	case token.AliasType:
  71  		if p.Anchor != nil {
  72  			return p.Alias()
  73  		}
  74  		return prop
  75  	case token.StringType, token.SingleQuoteType, token.DoubleQuoteType:
  76  		if p.String != nil {
  77  			return p.String()
  78  		}
  79  		return prop
  80  	case token.IntegerType, token.FloatType:
  81  		if p.Number != nil {
  82  			return p.Number()
  83  		}
  84  		return prop
  85  	default:
  86  	}
  87  	return prop
  88  }
  89  
  90  // PrintTokens create text from token collection
  91  func (p *Printer) PrintTokens(tokens token.Tokens) string {
  92  	if len(tokens) == 0 {
  93  		return ""
  94  	}
  95  	if p.LineNumber {
  96  		if p.LineNumberFormat == nil {
  97  			p.LineNumberFormat = defaultLineNumberFormat
  98  		}
  99  	}
 100  	texts := []string{}
 101  	lineNumber := tokens[0].Position.Line
 102  	for _, tk := range tokens {
 103  		lines := strings.Split(tk.Origin, "\n")
 104  		prop := p.property(tk)
 105  		header := ""
 106  		if p.LineNumber {
 107  			header = p.LineNumberFormat(lineNumber)
 108  		}
 109  		if len(lines) == 1 {
 110  			line := prop.Prefix + lines[0] + prop.Suffix
 111  			if len(texts) == 0 {
 112  				texts = append(texts, header+line)
 113  				lineNumber++
 114  			} else {
 115  				text := texts[len(texts)-1]
 116  				texts[len(texts)-1] = text + line
 117  			}
 118  		} else {
 119  			for idx, src := range lines {
 120  				if p.LineNumber {
 121  					header = p.LineNumberFormat(lineNumber)
 122  				}
 123  				line := prop.Prefix + src + prop.Suffix
 124  				if idx == 0 {
 125  					if len(texts) == 0 {
 126  						texts = append(texts, header+line)
 127  						lineNumber++
 128  					} else {
 129  						text := texts[len(texts)-1]
 130  						texts[len(texts)-1] = text + line
 131  					}
 132  				} else {
 133  					texts = append(texts, fmt.Sprintf("%s%s", header, line))
 134  					lineNumber++
 135  				}
 136  			}
 137  		}
 138  	}
 139  	return strings.Join(texts, "\n")
 140  }
 141  
 142  // PrintNode create text from ast.Node
 143  func (p *Printer) PrintNode(node ast.Node) []byte {
 144  	return []byte(fmt.Sprintf("%+v\n", node))
 145  }
 146  
 147  const escape = "\x1b"
 148  
 149  func format(attr color.Attribute) string {
 150  	return fmt.Sprintf("%s[%dm", escape, attr)
 151  }
 152  
 153  func (p *Printer) setDefaultColorSet() {
 154  	p.Bool = func() *Property {
 155  		return &Property{
 156  			Prefix: format(color.FgHiMagenta),
 157  			Suffix: format(color.Reset),
 158  		}
 159  	}
 160  	p.Number = func() *Property {
 161  		return &Property{
 162  			Prefix: format(color.FgHiMagenta),
 163  			Suffix: format(color.Reset),
 164  		}
 165  	}
 166  	p.MapKey = func() *Property {
 167  		return &Property{
 168  			Prefix: format(color.FgHiCyan),
 169  			Suffix: format(color.Reset),
 170  		}
 171  	}
 172  	p.Anchor = func() *Property {
 173  		return &Property{
 174  			Prefix: format(color.FgHiYellow),
 175  			Suffix: format(color.Reset),
 176  		}
 177  	}
 178  	p.Alias = func() *Property {
 179  		return &Property{
 180  			Prefix: format(color.FgHiYellow),
 181  			Suffix: format(color.Reset),
 182  		}
 183  	}
 184  	p.String = func() *Property {
 185  		return &Property{
 186  			Prefix: format(color.FgHiGreen),
 187  			Suffix: format(color.Reset),
 188  		}
 189  	}
 190  }
 191  
 192  func (p *Printer) PrintErrorMessage(msg string, isColored bool) string {
 193  	if isColored {
 194  		return fmt.Sprintf("%s%s%s",
 195  			format(color.FgHiRed),
 196  			msg,
 197  			format(color.Reset),
 198  		)
 199  	}
 200  	return msg
 201  }
 202  
 203  func (p *Printer) removeLeftSideNewLineChar(src string) string {
 204  	return strings.TrimLeft(strings.TrimLeft(strings.TrimLeft(src, "\r"), "\n"), "\r\n")
 205  }
 206  
 207  func (p *Printer) removeRightSideNewLineChar(src string) string {
 208  	return strings.TrimRight(strings.TrimRight(strings.TrimRight(src, "\r"), "\n"), "\r\n")
 209  }
 210  
 211  func (p *Printer) removeRightSideWhiteSpaceChar(src string) string {
 212  	return p.removeRightSideNewLineChar(strings.TrimRight(src, " "))
 213  }
 214  
 215  func (p *Printer) newLineCount(s string) int {
 216  	src := []rune(s)
 217  	size := len(src)
 218  	cnt := 0
 219  	for i := 0; i < size; i++ {
 220  		c := src[i]
 221  		switch c {
 222  		case '\r':
 223  			if i+1 < size && src[i+1] == '\n' {
 224  				i++
 225  			}
 226  			cnt++
 227  		case '\n':
 228  			cnt++
 229  		}
 230  	}
 231  	return cnt
 232  }
 233  
 234  func (p *Printer) isNewLineLastChar(s string) bool {
 235  	for i := len(s) - 1; i > 0; i-- {
 236  		c := s[i]
 237  		switch c {
 238  		case ' ':
 239  			continue
 240  		case '\n', '\r':
 241  			return true
 242  		}
 243  		break
 244  	}
 245  	return false
 246  }
 247  
 248  func (p *Printer) printBeforeTokens(tk *token.Token, minLine, extLine int) token.Tokens {
 249  	for {
 250  		if tk.Prev == nil {
 251  			break
 252  		}
 253  		if tk.Prev.Position.Line < minLine {
 254  			break
 255  		}
 256  		tk = tk.Prev
 257  	}
 258  	minTk := tk.Clone()
 259  	if minTk.Prev != nil {
 260  		// add white spaces to minTk by prev token
 261  		prev := minTk.Prev
 262  		whiteSpaceLen := len(prev.Origin) - len(strings.TrimRight(prev.Origin, " "))
 263  		minTk.Origin = strings.Repeat(" ", whiteSpaceLen) + minTk.Origin
 264  	}
 265  	minTk.Origin = p.removeLeftSideNewLineChar(minTk.Origin)
 266  	tokens := token.Tokens{minTk}
 267  	tk = minTk.Next
 268  	for tk != nil && tk.Position.Line <= extLine {
 269  		clonedTk := tk.Clone()
 270  		tokens.Add(clonedTk)
 271  		tk = clonedTk.Next
 272  	}
 273  	lastTk := tokens[len(tokens)-1]
 274  	trimmedOrigin := p.removeRightSideWhiteSpaceChar(lastTk.Origin)
 275  	suffix := lastTk.Origin[len(trimmedOrigin):]
 276  	lastTk.Origin = trimmedOrigin
 277  
 278  	if lastTk.Next != nil && len(suffix) > 1 {
 279  		next := lastTk.Next.Clone()
 280  		// add suffix to header of next token
 281  		if suffix[0] == '\n' || suffix[0] == '\r' {
 282  			suffix = suffix[1:]
 283  		}
 284  		next.Origin = suffix + next.Origin
 285  		lastTk.Next = next
 286  	}
 287  	return tokens
 288  }
 289  
 290  func (p *Printer) printAfterTokens(tk *token.Token, maxLine int) token.Tokens {
 291  	tokens := token.Tokens{}
 292  	if tk == nil {
 293  		return tokens
 294  	}
 295  	if tk.Position.Line > maxLine {
 296  		return tokens
 297  	}
 298  	minTk := tk.Clone()
 299  	minTk.Origin = p.removeLeftSideNewLineChar(minTk.Origin)
 300  	tokens.Add(minTk)
 301  	tk = minTk.Next
 302  	for tk != nil && tk.Position.Line <= maxLine {
 303  		clonedTk := tk.Clone()
 304  		tokens.Add(clonedTk)
 305  		tk = clonedTk.Next
 306  	}
 307  	return tokens
 308  }
 309  
 310  func (p *Printer) setupErrorTokenFormat(annotateLine int, isColored bool) {
 311  	prefix := func(annotateLine, num int) string {
 312  		if annotateLine == num {
 313  			return fmt.Sprintf("> %2d | ", num)
 314  		}
 315  		return fmt.Sprintf("  %2d | ", num)
 316  	}
 317  	p.LineNumber = true
 318  	p.LineNumberFormat = func(num int) string {
 319  		if isColored {
 320  			fn := color.New(color.Bold, color.FgHiWhite).SprintFunc()
 321  			return fn(prefix(annotateLine, num))
 322  		}
 323  		return prefix(annotateLine, num)
 324  	}
 325  	if isColored {
 326  		p.setDefaultColorSet()
 327  	}
 328  }
 329  
 330  func (p *Printer) PrintErrorToken(tk *token.Token, isColored bool) string {
 331  	errToken := tk
 332  	curLine := tk.Position.Line
 333  	curExtLine := curLine + p.newLineCount(p.removeLeftSideNewLineChar(tk.Origin))
 334  	if p.isNewLineLastChar(tk.Origin) {
 335  		// if last character ( exclude white space ) is new line character, ignore it.
 336  		curExtLine--
 337  	}
 338  
 339  	minLine := int(math.Max(float64(curLine-3), 1))
 340  	maxLine := curExtLine + 3
 341  	p.setupErrorTokenFormat(curLine, isColored)
 342  
 343  	beforeTokens := p.printBeforeTokens(tk, minLine, curExtLine)
 344  	lastTk := beforeTokens[len(beforeTokens)-1]
 345  	afterTokens := p.printAfterTokens(lastTk.Next, maxLine)
 346  
 347  	beforeSource := p.PrintTokens(beforeTokens)
 348  	prefixSpaceNum := len(fmt.Sprintf("  %2d | ", curLine))
 349  	annotateLine := strings.Repeat(" ", prefixSpaceNum+errToken.Position.Column-1) + "^"
 350  	afterSource := p.PrintTokens(afterTokens)
 351  	return fmt.Sprintf("%s\n%s\n%s", beforeSource, annotateLine, afterSource)
 352  }
 353