1-- Easy way to print data structes
2-- From https://github.com/jagt/pprint.lua, file is license as pubic domain
3
4local pprint = { VERSION = '0.1' }
5
6pprint.defaults = {
7    -- type display trigger, hide not useful datatypes by default
8    -- custom types are treated as table
9    show_nil = true,
10    show_boolean = true,
11    show_number = true,
12    show_string = true,
13    show_table = true,
14    show_function = false,
15    show_thread = false,
16    show_userdata = false,
17    -- additional display trigger
18    show_metatable = false,     -- show metatable
19    show_all = false,           -- override other show settings and show everything
20    use_tostring = false,       -- use __tostring to print table if available
21    filter_function = nil,      -- called like callback(value[,key, parent]), return truty value to hide
22    object_cache = 'local',     -- cache blob and table to give it a id, 'local' cache per print, 'global' cache
23                                -- per process, falsy value to disable (might cause infinite loop)
24    -- format settings
25    indent_size = 2,            -- indent for each nested table level
26    level_width = 80,           -- max width per indent level
27    wrap_string = true,         -- wrap string when it's longer than level_width
28    wrap_array = false,         -- wrap every array elements
29    sort_keys = true,           -- sort table keys
30}
31
32local TYPES = {
33    ['nil'] = 1, ['boolean'] = 2, ['number'] = 3, ['string'] = 4,
34    ['table'] = 5, ['function'] = 6, ['thread'] = 7, ['userdata'] = 8
35}
36
37-- seems this is the only way to escape these, as lua don't know how to map char '\a' to 'a'
38local ESCAPE_MAP = {
39    ['\a'] = '\\a', ['\b'] = '\\b', ['\f'] = '\\f', ['\n'] = '\\n', ['\r'] = '\\r',
40    ['\t'] = '\\t', ['\v'] = '\\v', ['\\'] = '\\\\',
41}
42
43-- generic utilities
44local function escape(s)
45    s = s:gsub('([%c\\])', ESCAPE_MAP)
46    local dq = s:find('"')
47    local sq = s:find("'")
48    if dq and sq then
49        return s:gsub('"', '\\"'), '"'
50    elseif sq then
51        return s, '"'
52    else
53        return s, "'"
54    end
55end
56
57local function is_plain_key(key)
58    return type(key) == 'string' and key:match('^[%a_][%a%d_]*$')
59end
60
61local CACHE_TYPES = {
62    ['table'] = true, ['function'] = true, ['thread'] = true, ['userdata'] = true
63}
64
65-- cache would be populated to be like:
66-- {
67--     function = { `fun1` = 1, _cnt = 1 }, -- object id
68--     table = { `table1` = 1, `table2` = 2, _cnt = 2 },
69--     visited_tables = { `table1` = 7, `table2` = 8  }, -- visit count
70-- }
71-- use weakrefs to avoid accidentall adding refcount
72local function cache_apperance(obj, cache, option)
73    if not cache.visited_tables then
74        cache.visited_tables = setmetatable({}, {__mode = 'k'})
75    end
76    local t = type(obj)
77
78    -- TODO can't test filter_function here as we don't have the ix and key,
79    -- might cause different results?
80    -- respect show_xxx and filter_function to be consistent with print results
81    if (not TYPES[t] and not option.show_table)
82        or (TYPES[t] and not option['show_'..t]) then
83        return
84    end
85
86    if CACHE_TYPES[t] or TYPES[t] == nil then
87        if not cache[t] then
88            cache[t] = setmetatable({}, {__mode = 'k'})
89            cache[t]._cnt = 0
90        end
91        if not cache[t][obj] then
92            cache[t]._cnt = cache[t]._cnt + 1
93            cache[t][obj] = cache[t]._cnt
94        end
95    end
96    if t == 'table' or TYPES[t] == nil then
97        if cache.visited_tables[obj] == false then
98            -- already printed, no need to mark this and its children anymore
99            return
100        elseif cache.visited_tables[obj] == nil then
101            cache.visited_tables[obj] = 1
102        else
103            -- visited already, increment and continue
104            cache.visited_tables[obj] = cache.visited_tables[obj] + 1
105            return
106        end
107        for k, v in pairs(obj) do
108            cache_apperance(k, cache, option)
109            cache_apperance(v, cache, option)
110        end
111        local mt = getmetatable(obj)
112        if mt and option.show_metatable then
113            cache_apperance(mt, cache, option)
114        end
115    end
116end
117
118-- makes 'foo2' < 'foo100000'. string.sub makes substring anyway, no need to use index based method
119local function str_natural_cmp(lhs, rhs)
120    while #lhs > 0 and #rhs > 0 do
121        local lmid, lend = lhs:find('%d+')
122        local rmid, rend = rhs:find('%d+')
123        if not (lmid and rmid) then return lhs < rhs end
124
125        local lsub = lhs:sub(1, lmid-1)
126        local rsub = rhs:sub(1, rmid-1)
127        if lsub ~= rsub then
128            return lsub < rsub
129        end
130
131        local lnum = tonumber(lhs:sub(lmid, lend))
132        local rnum = tonumber(rhs:sub(rmid, rend))
133        if lnum ~= rnum then
134            return lnum < rnum
135        end
136
137        lhs = lhs:sub(lend+1)
138        rhs = rhs:sub(rend+1)
139    end
140    return lhs < rhs
141end
142
143local function cmp(lhs, rhs)
144    local tleft = type(lhs)
145    local tright = type(rhs)
146    if tleft == 'number' and tright == 'number' then return lhs < rhs end
147    if tleft == 'string' and tright == 'string' then return str_natural_cmp(lhs, rhs) end
148    if tleft == tright then return str_natural_cmp(tostring(lhs), tostring(rhs)) end
149
150    -- allow custom types
151    local oleft = TYPES[tleft] or 9
152    local oright = TYPES[tright] or 9
153    return oleft < oright
154end
155
156-- setup option with default
157local function make_option(option)
158    if option == nil then
159        option = {}
160    end
161    for k, v in pairs(pprint.defaults) do
162        if option[k] == nil then
163            option[k] = v
164        end
165        if option.show_all then
166            for t, _ in pairs(TYPES) do
167                option['show_'..t] = true
168            end
169            option.show_metatable = true
170        end
171    end
172    return option
173end
174
175-- override defaults and take effects for all following calls
176function pprint.setup(option)
177    pprint.defaults = make_option(option)
178end
179
180-- format lua object into a string
181function pprint.pformat(obj, option, printer)
182    option = make_option(option)
183    local buf = {}
184    local function default_printer(s)
185        table.insert(buf, s)
186    end
187    printer = printer or default_printer
188
189    local cache
190    if option.object_cache == 'global' then
191        -- steal the cache into a local var so it's not visible from _G or anywhere
192        -- still can't avoid user explicitly referentce pprint._cache but it shouldn't happen anyway
193        cache = pprint._cache or {}
194        pprint._cache = nil
195    elseif option.object_cache == 'local' then
196        cache = {}
197    end
198
199    local last = '' -- used for look back and remove trailing comma
200    local status = {
201        indent = '', -- current indent
202        len = 0,     -- current line length
203    }
204
205    local wrapped_printer = function(s)
206        printer(last)
207        last = s
208    end
209
210    local function _indent(d)
211        status.indent = string.rep(' ', d + #(status.indent))
212    end
213
214    local function _n(d)
215        wrapped_printer('\n')
216        wrapped_printer(status.indent)
217        if d then
218            _indent(d)
219        end
220        status.len = 0
221        return true -- used to close bracket correctly
222    end
223
224    local function _p(s, nowrap)
225        status.len = status.len + #s
226        if not nowrap and status.len > option.level_width then
227            _n()
228            wrapped_printer(s)
229            status.len = #s
230        else
231            wrapped_printer(s)
232        end
233    end
234
235    local formatter = {}
236    local function format(v)
237        local f = formatter[type(v)]
238        f = f or formatter.table -- allow patched type()
239        if option.filter_function and option.filter_function(v, nil, nil) then
240            return ''
241        else
242            return f(v)
243        end
244    end
245
246    local function tostring_formatter(v)
247        return tostring(v)
248    end
249
250    local function number_formatter(n)
251        return n == math.huge and '[[math.huge]]' or tostring(n)
252    end
253
254    local function nop_formatter(v)
255        return ''
256    end
257
258    local function make_fixed_formatter(t, has_cache)
259        if has_cache then
260            return function (v)
261                return string.format('[[%s %d]]', t, cache[t][v])
262            end
263        else
264            return function (v)
265                return '[['..t..']]'
266            end
267        end
268    end
269
270    local function string_formatter(s, force_long_quote)
271        local s, quote = escape(s)
272        local quote_len = force_long_quote and 4 or 2
273        if quote_len + #s + status.len > option.level_width then
274            _n()
275            -- only wrap string when is longer than level_width
276            if option.wrap_string and #s + quote_len > option.level_width then
277                -- keep the quotes together
278                _p('[[')
279                while #s + status.len >= option.level_width do
280                    local seg = option.level_width - status.len
281                    _p(string.sub(s, 1, seg), true)
282                    _n()
283                    s = string.sub(s, seg+1)
284                end
285                _p(s) -- print the remaining parts
286                return ']]'
287            end
288        end
289
290        return force_long_quote and '[['..s..']]' or quote..s..quote
291    end
292
293    local function table_formatter(t)
294        if option.use_tostring then
295            local mt = getmetatable(t)
296            if mt and mt.__tostring then
297                return string_formatter(tostring(t), true)
298            end
299        end
300
301        local print_header_ix = nil
302        local ttype = type(t)
303        if option.object_cache then
304            local cache_state = cache.visited_tables[t]
305            local tix = cache[ttype][t]
306            -- FIXME should really handle `cache_state == nil`
307            -- as user might add things through filter_function
308            if cache_state == false then
309                -- already printed, just print the the number
310                return string_formatter(string.format('%s %d', ttype, tix), true)
311            elseif cache_state > 1 then
312                -- appeared more than once, print table header with number
313                print_header_ix = tix
314                cache.visited_tables[t] = false
315            else
316                -- appeared exactly once, print like a normal table
317            end
318        end
319
320        local tlen = #t
321        local wrapped = false
322        _p('{')
323        _indent(option.indent_size)
324        _p(string.rep(' ', option.indent_size - 1))
325        if print_header_ix then
326            _p(string.format('--[[%s %d]] ', ttype, print_header_ix))
327        end
328        for ix = 1,tlen do
329            local v = t[ix]
330            if formatter[type(v)] == nop_formatter or
331               (option.filter_function and option.filter_function(v, ix, t)) then
332               -- pass
333            else
334                if option.wrap_array then
335                    wrapped = _n()
336                end
337                _p(format(v)..', ')
338            end
339        end
340
341        -- hashmap part of the table, in contrast to array part
342        local function is_hash_key(k)
343            local numkey = tonumber(k)
344            if numkey ~= k or numkey > tlen then
345                return true
346            end
347        end
348
349        local function print_kv(k, v, t)
350            -- can't use option.show_x as obj may contain custom type
351            if formatter[type(v)] == nop_formatter or
352               formatter[type(k)] == nop_formatter or
353               (option.filter_function and option.filter_function(v, k, t)) then
354                return
355            end
356            wrapped = _n()
357            if is_plain_key(k) then
358                _p(k, true)
359            else
360                _p('[')
361                -- [[]] type string in key is illegal, needs to add spaces inbetween
362                local k = format(k)
363                if string.match(k, '%[%[') then
364                    _p(' '..k..' ', true)
365                else
366                    _p(k, true)
367                end
368                _p(']')
369            end
370            _p(' = ', true)
371            _p(format(v), true)
372            _p(',', true)
373        end
374
375        if option.sort_keys then
376            local keys = {}
377            for k, _ in pairs(t) do
378                if is_hash_key(k) then
379                    table.insert(keys, k)
380                end
381            end
382            table.sort(keys, cmp)
383            for _, k in ipairs(keys) do
384                print_kv(k, t[k], t)
385            end
386        else
387            for k, v in pairs(t) do
388                if is_hash_key(k) then
389                    print_kv(k, v, t)
390                end
391            end
392        end
393
394        if option.show_metatable then
395            local mt = getmetatable(t)
396            if mt then
397                print_kv('__metatable', mt, t)
398            end
399        end
400
401        _indent(-option.indent_size)
402        -- make { } into {}
403        last = string.gsub(last, '^ +$', '')
404        -- peek last to remove trailing comma
405        last = string.gsub(last, ',%s*$', ' ')
406        if wrapped then
407            _n()
408        end
409        _p('}')
410
411        return ''
412    end
413
414    -- set formatters
415    formatter['nil'] = option.show_nil and tostring_formatter or nop_formatter
416    formatter['boolean'] = option.show_boolean and tostring_formatter or nop_formatter
417    formatter['number'] = option.show_number and number_formatter or nop_formatter -- need to handle math.huge
418    formatter['function'] = option.show_function and make_fixed_formatter('function', option.object_cache) or nop_formatter
419    formatter['thread'] = option.show_thread and make_fixed_formatter('thread', option.object_cache) or nop_formatter
420    formatter['userdata'] = option.show_userdata and make_fixed_formatter('userdata', option.object_cache) or nop_formatter
421    formatter['string'] = option.show_string and string_formatter or nop_formatter
422    formatter['table'] = option.show_table and table_formatter or nop_formatter
423
424    if option.object_cache then
425        -- needs to visit the table before start printing
426        cache_apperance(obj, cache, option)
427    end
428
429    _p(format(obj))
430    printer(last) -- close the buffered one
431
432    -- put cache back if global
433    if option.object_cache == 'global' then
434        pprint._cache = cache
435    end
436
437    return table.concat(buf)
438end
439
440-- pprint all the arguments
441function pprint.pprint( ... )
442    local args = {...}
443    -- select will get an accurate count of array len, counting trailing nils
444    local len = select('#', ...)
445    for ix = 1,len do
446        pprint.pformat(args[ix], nil, io.write)
447        io.write('\n')
448    end
449end
450
451setmetatable(pprint, {
452    __call = function (_, ...)
453        pprint.pprint(...)
454    end
455})
456
457return pprint
458