1package pretty
2
3import (
4	"sort"
5)
6
7// Options is Pretty options
8type Options struct {
9	// Width is an max column width for single line arrays
10	// Default is 80
11	Width int
12	// Prefix is a prefix for all lines
13	// Default is an empty string
14	Prefix string
15	// Indent is the nested indentation
16	// Default is two spaces
17	Indent string
18	// SortKeys will sort the keys alphabetically
19	// Default is false
20	SortKeys bool
21}
22
23// DefaultOptions is the default options for pretty formats.
24var DefaultOptions = &Options{Width: 80, Prefix: "", Indent: "  ", SortKeys: false}
25
26// Pretty converts the input json into a more human readable format where each
27// element is on it's own line with clear indentation.
28func Pretty(json []byte) []byte { return PrettyOptions(json, nil) }
29
30// PrettyOptions is like Pretty but with customized options.
31func PrettyOptions(json []byte, opts *Options) []byte {
32	if opts == nil {
33		opts = DefaultOptions
34	}
35	buf := make([]byte, 0, len(json))
36	if len(opts.Prefix) != 0 {
37		buf = append(buf, opts.Prefix...)
38	}
39	buf, _, _, _ = appendPrettyAny(buf, json, 0, true,
40		opts.Width, opts.Prefix, opts.Indent, opts.SortKeys,
41		0, 0, -1)
42	if len(buf) > 0 {
43		buf = append(buf, '\n')
44	}
45	return buf
46}
47
48// Ugly removes insignificant space characters from the input json byte slice
49// and returns the compacted result.
50func Ugly(json []byte) []byte {
51	buf := make([]byte, 0, len(json))
52	return ugly(buf, json)
53}
54
55// UglyInPlace removes insignificant space characters from the input json
56// byte slice and returns the compacted result. This method reuses the
57// input json buffer to avoid allocations. Do not use the original bytes
58// slice upon return.
59func UglyInPlace(json []byte) []byte { return ugly(json, json) }
60
61func ugly(dst, src []byte) []byte {
62	dst = dst[:0]
63	for i := 0; i < len(src); i++ {
64		if src[i] > ' ' {
65			dst = append(dst, src[i])
66			if src[i] == '"' {
67				for i = i + 1; i < len(src); i++ {
68					dst = append(dst, src[i])
69					if src[i] == '"' {
70						j := i - 1
71						for ; ; j-- {
72							if src[j] != '\\' {
73								break
74							}
75						}
76						if (j-i)%2 != 0 {
77							break
78						}
79					}
80				}
81			}
82		}
83	}
84	return dst
85}
86
87func appendPrettyAny(buf, json []byte, i int, pretty bool, width int, prefix, indent string, sortkeys bool, tabs, nl, max int) ([]byte, int, int, bool) {
88	for ; i < len(json); i++ {
89		if json[i] <= ' ' {
90			continue
91		}
92		if json[i] == '"' {
93			return appendPrettyString(buf, json, i, nl)
94		}
95		if (json[i] >= '0' && json[i] <= '9') || json[i] == '-' {
96			return appendPrettyNumber(buf, json, i, nl)
97		}
98		if json[i] == '{' {
99			return appendPrettyObject(buf, json, i, '{', '}', pretty, width, prefix, indent, sortkeys, tabs, nl, max)
100		}
101		if json[i] == '[' {
102			return appendPrettyObject(buf, json, i, '[', ']', pretty, width, prefix, indent, sortkeys, tabs, nl, max)
103		}
104		switch json[i] {
105		case 't':
106			return append(buf, 't', 'r', 'u', 'e'), i + 4, nl, true
107		case 'f':
108			return append(buf, 'f', 'a', 'l', 's', 'e'), i + 5, nl, true
109		case 'n':
110			return append(buf, 'n', 'u', 'l', 'l'), i + 4, nl, true
111		}
112	}
113	return buf, i, nl, true
114}
115
116type pair struct {
117	kstart, kend int
118	vstart, vend int
119}
120
121type byKey struct {
122	sorted bool
123	json   []byte
124	pairs  []pair
125}
126
127func (arr *byKey) Len() int {
128	return len(arr.pairs)
129}
130func (arr *byKey) Less(i, j int) bool {
131	key1 := arr.json[arr.pairs[i].kstart+1 : arr.pairs[i].kend-1]
132	key2 := arr.json[arr.pairs[j].kstart+1 : arr.pairs[j].kend-1]
133	return string(key1) < string(key2)
134}
135func (arr *byKey) Swap(i, j int) {
136	arr.pairs[i], arr.pairs[j] = arr.pairs[j], arr.pairs[i]
137	arr.sorted = true
138}
139
140func appendPrettyObject(buf, json []byte, i int, open, close byte, pretty bool, width int, prefix, indent string, sortkeys bool, tabs, nl, max int) ([]byte, int, int, bool) {
141	var ok bool
142	if width > 0 {
143		if pretty && open == '[' && max == -1 {
144			// here we try to create a single line array
145			max := width - (len(buf) - nl)
146			if max > 3 {
147				s1, s2 := len(buf), i
148				buf, i, _, ok = appendPrettyObject(buf, json, i, '[', ']', false, width, prefix, "", sortkeys, 0, 0, max)
149				if ok && len(buf)-s1 <= max {
150					return buf, i, nl, true
151				}
152				buf = buf[:s1]
153				i = s2
154			}
155		} else if max != -1 && open == '{' {
156			return buf, i, nl, false
157		}
158	}
159	buf = append(buf, open)
160	i++
161	var pairs []pair
162	if open == '{' && sortkeys {
163		pairs = make([]pair, 0, 8)
164	}
165	var n int
166	for ; i < len(json); i++ {
167		if json[i] <= ' ' {
168			continue
169		}
170		if json[i] == close {
171			if pretty {
172				if open == '{' && sortkeys {
173					buf = sortPairs(json, buf, pairs)
174				}
175				if n > 0 {
176					nl = len(buf)
177					buf = append(buf, '\n')
178				}
179				if buf[len(buf)-1] != open {
180					buf = appendTabs(buf, prefix, indent, tabs)
181				}
182			}
183			buf = append(buf, close)
184			return buf, i + 1, nl, open != '{'
185		}
186		if open == '[' || json[i] == '"' {
187			if n > 0 {
188				buf = append(buf, ',')
189				if width != -1 && open == '[' {
190					buf = append(buf, ' ')
191				}
192			}
193			var p pair
194			if pretty {
195				nl = len(buf)
196				buf = append(buf, '\n')
197				if open == '{' && sortkeys {
198					p.kstart = i
199					p.vstart = len(buf)
200				}
201				buf = appendTabs(buf, prefix, indent, tabs+1)
202			}
203			if open == '{' {
204				buf, i, nl, _ = appendPrettyString(buf, json, i, nl)
205				if sortkeys {
206					p.kend = i
207				}
208				buf = append(buf, ':')
209				if pretty {
210					buf = append(buf, ' ')
211				}
212			}
213			buf, i, nl, ok = appendPrettyAny(buf, json, i, pretty, width, prefix, indent, sortkeys, tabs+1, nl, max)
214			if max != -1 && !ok {
215				return buf, i, nl, false
216			}
217			if pretty && open == '{' && sortkeys {
218				p.vend = len(buf)
219				if p.kstart > p.kend || p.vstart > p.vend {
220					// bad data. disable sorting
221					sortkeys = false
222				} else {
223					pairs = append(pairs, p)
224				}
225			}
226			i--
227			n++
228		}
229	}
230	return buf, i, nl, open != '{'
231}
232func sortPairs(json, buf []byte, pairs []pair) []byte {
233	if len(pairs) == 0 {
234		return buf
235	}
236	vstart := pairs[0].vstart
237	vend := pairs[len(pairs)-1].vend
238	arr := byKey{false, json, pairs}
239	sort.Sort(&arr)
240	if !arr.sorted {
241		return buf
242	}
243	nbuf := make([]byte, 0, vend-vstart)
244	for i, p := range pairs {
245		nbuf = append(nbuf, buf[p.vstart:p.vend]...)
246		if i < len(pairs)-1 {
247			nbuf = append(nbuf, ',')
248			nbuf = append(nbuf, '\n')
249		}
250	}
251	return append(buf[:vstart], nbuf...)
252}
253
254func appendPrettyString(buf, json []byte, i, nl int) ([]byte, int, int, bool) {
255	s := i
256	i++
257	for ; i < len(json); i++ {
258		if json[i] == '"' {
259			var sc int
260			for j := i - 1; j > s; j-- {
261				if json[j] == '\\' {
262					sc++
263				} else {
264					break
265				}
266			}
267			if sc%2 == 1 {
268				continue
269			}
270			i++
271			break
272		}
273	}
274	return append(buf, json[s:i]...), i, nl, true
275}
276
277func appendPrettyNumber(buf, json []byte, i, nl int) ([]byte, int, int, bool) {
278	s := i
279	i++
280	for ; i < len(json); i++ {
281		if json[i] <= ' ' || json[i] == ',' || json[i] == ':' || json[i] == ']' || json[i] == '}' {
282			break
283		}
284	}
285	return append(buf, json[s:i]...), i, nl, true
286}
287
288func appendTabs(buf []byte, prefix, indent string, tabs int) []byte {
289	if len(prefix) != 0 {
290		buf = append(buf, prefix...)
291	}
292	if len(indent) == 2 && indent[0] == ' ' && indent[1] == ' ' {
293		for i := 0; i < tabs; i++ {
294			buf = append(buf, ' ', ' ')
295		}
296	} else {
297		for i := 0; i < tabs; i++ {
298			buf = append(buf, indent...)
299		}
300	}
301	return buf
302}
303
304// Style is the color style
305type Style struct {
306	Key, String, Number [2]string
307	True, False, Null   [2]string
308	Append              func(dst []byte, c byte) []byte
309}
310
311func hexp(p byte) byte {
312	switch {
313	case p < 10:
314		return p + '0'
315	default:
316		return (p - 10) + 'a'
317	}
318}
319
320// TerminalStyle is for terminals
321var TerminalStyle = &Style{
322	Key:    [2]string{"\x1B[94m", "\x1B[0m"},
323	String: [2]string{"\x1B[92m", "\x1B[0m"},
324	Number: [2]string{"\x1B[93m", "\x1B[0m"},
325	True:   [2]string{"\x1B[96m", "\x1B[0m"},
326	False:  [2]string{"\x1B[96m", "\x1B[0m"},
327	Null:   [2]string{"\x1B[91m", "\x1B[0m"},
328	Append: func(dst []byte, c byte) []byte {
329		if c < ' ' && (c != '\r' && c != '\n' && c != '\t' && c != '\v') {
330			dst = append(dst, "\\u00"...)
331			dst = append(dst, hexp((c>>4)&0xF))
332			return append(dst, hexp((c)&0xF))
333		}
334		return append(dst, c)
335	},
336}
337
338// Color will colorize the json. The style parma is used for customizing
339// the colors. Passing nil to the style param will use the default
340// TerminalStyle.
341func Color(src []byte, style *Style) []byte {
342	if style == nil {
343		style = TerminalStyle
344	}
345	apnd := style.Append
346	if apnd == nil {
347		apnd = func(dst []byte, c byte) []byte {
348			return append(dst, c)
349		}
350	}
351	type stackt struct {
352		kind byte
353		key  bool
354	}
355	var dst []byte
356	var stack []stackt
357	for i := 0; i < len(src); i++ {
358		if src[i] == '"' {
359			key := len(stack) > 0 && stack[len(stack)-1].key
360			if key {
361				dst = append(dst, style.Key[0]...)
362			} else {
363				dst = append(dst, style.String[0]...)
364			}
365			dst = apnd(dst, '"')
366			for i = i + 1; i < len(src); i++ {
367				dst = apnd(dst, src[i])
368				if src[i] == '"' {
369					j := i - 1
370					for ; ; j-- {
371						if src[j] != '\\' {
372							break
373						}
374					}
375					if (j-i)%2 != 0 {
376						break
377					}
378				}
379			}
380			if key {
381				dst = append(dst, style.Key[1]...)
382			} else {
383				dst = append(dst, style.String[1]...)
384			}
385		} else if src[i] == '{' || src[i] == '[' {
386			stack = append(stack, stackt{src[i], src[i] == '{'})
387			dst = apnd(dst, src[i])
388		} else if (src[i] == '}' || src[i] == ']') && len(stack) > 0 {
389			stack = stack[:len(stack)-1]
390			dst = apnd(dst, src[i])
391		} else if (src[i] == ':' || src[i] == ',') && len(stack) > 0 && stack[len(stack)-1].kind == '{' {
392			stack[len(stack)-1].key = !stack[len(stack)-1].key
393			dst = apnd(dst, src[i])
394		} else {
395			var kind byte
396			if (src[i] >= '0' && src[i] <= '9') || src[i] == '-' {
397				kind = '0'
398				dst = append(dst, style.Number[0]...)
399			} else if src[i] == 't' {
400				kind = 't'
401				dst = append(dst, style.True[0]...)
402			} else if src[i] == 'f' {
403				kind = 'f'
404				dst = append(dst, style.False[0]...)
405			} else if src[i] == 'n' {
406				kind = 'n'
407				dst = append(dst, style.Null[0]...)
408			} else {
409				dst = apnd(dst, src[i])
410			}
411			if kind != 0 {
412				for ; i < len(src); i++ {
413					if src[i] <= ' ' || src[i] == ',' || src[i] == ':' || src[i] == ']' || src[i] == '}' {
414						i--
415						break
416					}
417					dst = apnd(dst, src[i])
418				}
419				if kind == '0' {
420					dst = append(dst, style.Number[1]...)
421				} else if kind == 't' {
422					dst = append(dst, style.True[1]...)
423				} else if kind == 'f' {
424					dst = append(dst, style.False[1]...)
425				} else if kind == 'n' {
426					dst = append(dst, style.Null[1]...)
427				}
428			}
429		}
430	}
431	return dst
432}
433