决策树可视化python代码

#这段代码直接将决策树突出的模型转成json类型

from __future__ import print_function

import json

if __name__ == "__main__":

    treetest = open("data/testif.txt").read()
    def parse(lines):
        block = []
        while lines:
            if lines[0].startswith('If'):
                bl = ' '.join(lines.pop(0).split()[1:]).replace('(', '').replace(')', '')
                print("bl"+bl)
                block.append({
 
 'name': bl, 'children': parse(lines)})
                if lines[0].startswith('Else'):
                    be = ' '.join(lines.pop(0).split()[1:]).replace('(', '').replace(')', '')
                    block.append({
 
 'name': be, 'children': parse(lines)})
            elif not lines[0].startswith(('If', 'Else')):
                block2 = lines.pop(0)
                block.append({
 
 'name': block2})
            else:
                break
        return block

def tree_json(tree):
    data = []
    for line in tree.splitlines():
        if line.strip():
            line = line.strip()
            print(line+"*************")
            data.append(line)
            data2=data[1:]
            print(data2)
        else:
            break
        if not line: break
    res = []
    res.append({
 
 'name': 'Root', 'children': parse(data[0:])})
    with open('data/structure.json', 'w') as outfile:
        json.dump(res[0], outfile)
    print('Conversion Success !')
tree_json(treetest)