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
323func init() {
324	TerminalStyle = &Style{
325		Key:    [2]string{"\x1B[94m", "\x1B[0m"},
326		String: [2]string{"\x1B[92m", "\x1B[0m"},
327		Number: [2]string{"\x1B[93m", "\x1B[0m"},
328		True:   [2]string{"\x1B[96m", "\x1B[0m"},
329		False:  [2]string{"\x1B[96m", "\x1B[0m"},
330		Null:   [2]string{"\x1B[91m", "\x1B[0m"},
331		Append: func(dst []byte, c byte) []byte {
332			if c < ' ' && (c != '\r' && c != '\n' && c != '\t' && c != '\v') {
333				dst = append(dst, "\\u00"...)
334				dst = append(dst, hexp((c>>4)&0xF))
335				return append(dst, hexp((c)&0xF))
336			}
337			return append(dst, c)
338		},
339	}
340}
341
342// Color will colorize the json. The style parma is used for customizing
343// the colors. Passing nil to the style param will use the default
344// TerminalStyle.
345func Color(src []byte, style *Style) []byte {
346	if style == nil {
347		style = TerminalStyle
348	}
349	apnd := style.Append
350	if apnd == nil {
351		apnd = func(dst []byte, c byte) []byte {
352			return append(dst, c)
353		}
354	}
355	type stackt struct {
356		kind byte
357		key  bool
358	}
359	var dst []byte
360	var stack []stackt
361	for i := 0; i < len(src); i++ {
362		if src[i] == '"' {
363			key := len(stack) > 0 && stack[len(stack)-1].key
364			if key {
365				dst = append(dst, style.Key[0]...)
366			} else {
367				dst = append(dst, style.String[0]...)
368			}
369			dst = apnd(dst, '"')
370			for i = i + 1; i < len(src); i++ {
371				dst = apnd(dst, src[i])
372				if src[i] == '"' {
373					j := i - 1
374					for ; ; j-- {
375						if src[j] != '\\' {
376							break
377						}
378					}
379					if (j-i)%2 != 0 {
380						break
381					}
382				}
383			}
384			if key {
385				dst = append(dst, style.Key[1]...)
386			} else {
387				dst = append(dst, style.String[1]...)
388			}
389		} else if src[i] == '{' || src[i] == '[' {
390			stack = append(stack, stackt{src[i], src[i] == '{'})
391			dst = apnd(dst, src[i])
392		} else if (src[i] == '}' || src[i] == ']') && len(stack) > 0 {
393			stack = stack[:len(stack)-1]
394			dst = apnd(dst, src[i])
395		} else if (src[i] == ':' || src[i] == ',') && len(stack) > 0 && stack[len(stack)-1].kind == '{' {
396			stack[len(stack)-1].key = !stack[len(stack)-1].key
397			dst = apnd(dst, src[i])
398		} else {
399			var kind byte
400			if (src[i] >= '0' && src[i] <= '9') || src[i] == '-' {
401				kind = '0'
402				dst = append(dst, style.Number[0]...)
403			} else if src[i] == 't' {
404				kind = 't'
405				dst = append(dst, style.True[0]...)
406			} else if src[i] == 'f' {
407				kind = 'f'
408				dst = append(dst, style.False[0]...)
409			} else if src[i] == 'n' {
410				kind = 'n'
411				dst = append(dst, style.Null[0]...)
412			} else {
413				dst = apnd(dst, src[i])
414			}
415			if kind != 0 {
416				for ; i < len(src); i++ {
417					if src[i] <= ' ' || src[i] == ',' || src[i] == ':' || src[i] == ']' || src[i] == '}' {
418						i--
419						break
420					}
421					dst = apnd(dst, src[i])
422				}
423				if kind == '0' {
424					dst = append(dst, style.Number[1]...)
425				} else if kind == 't' {
426					dst = append(dst, style.True[1]...)
427				} else if kind == 'f' {
428					dst = append(dst, style.False[1]...)
429				} else if kind == 'n' {
430					dst = append(dst, style.Null[1]...)
431				}
432			}
433		}
434	}
435	return dst
436}
437