pretty print for xgboost gbtree booster dump

xgbdump_formatter.py

import sys
import re
import unicodedata

depth={}
yn={}

f1=None

f=open(sys.argv[1])
while True:
  line=f.readline().strip()
  if not line: break
  line = line.decode('sjis',"ignore")
  line = unicodedata.normalize('NFKC',unicode(line))

  m=re.match("booster\[([0-9]+)\]",line)
  if m:
    tree_id=int(m.group(1))+1
    depth={}
    yn={}
    if f1: f1.close()
    f1 = open("tree"+str(tree_id)+".txt","w")
  else:
    m=re.match("([0-9]+):",line)
    if m:
      id=m.group(1)
      if id not in depth:
	depth[id]=''
	yn[id]=False
      f1.write((depth[id]+line+"\n").encode("sjis"))
      m=re.match("[0-9]+:.* yes=([0-9]+),no=([0-9]+)",line)
      if m:
        yes=m.group(1)
        no=m.group(2)
        depth[yes]=depth[no]=depth[id]+('| ' if yn[id] else '  ')
        yn[yes]=True
        yn[no]=False
  
f.close()

if f1: f1.close()