pythonでcolor文字をprettyに出力

概要

colorでpretty print する記事が出てて、pythonではどうするんだと思い、stackoverflowとで部品をあさり、見よう見まねで書いてみた。

とりあえず、yamlファイルを読み込んでparseした結果を表示するようにした。

使い方

python pp.py <ファイル名>  [インデント数]

こんな感じ

pp.py

# -*- coding: utf-8 -*-

class color:
    BLUE = '\033[1;34m'
    BLUEL = '\033[0;34m'
    GREEN = '\033[1;32m'
    GREENL = '\033[0;32m'
    CYAN = '\033[1;36m'
    CYANL = '\033[0;36m'
    RED = '\033[1;31m'
    REDL = '\033[0;31m'
    PURPLE = '\033[1;35m'
    PURPLEL = '\033[0;35m'
    YELLOW = '\033[1;33m'
    BROWN = '\033[0;33m'
    WHITE = '\033[1;37m'
    GRAYL = '\033[0;37m'
    GRAYD = '\033[1;30m'
    BLACK = '\033[0;30m'
    ENDC = '\033[0m'

def pretty_convert(obj):
    if isinstance(obj, (float)):
        return u'{1}{2}{0}'.format(color.ENDC, color.PURPLEL, obj)
    elif isinstance(obj, ( bool )):
        return u'{1}{2}{0}'.format(color.ENDC, color.CYAN, obj)
    elif isinstance(obj, ( int, long, complex )):
        return u'{1}{2}{0}'.format(color.ENDC, color.BLUE, obj)
    elif isinstance(obj, ( basestring )):
        return u'{1}"{0}{2}{3}{0}{1}"{0}'.format(color.ENDC, color.REDL, color.RED, obj)
    elif isinstance(obj, ( dict )):
        return dict((u'{1}{2}:{0}'.format(color.ENDC, color.YELLOW,k), pretty_convert(v)) for k, v in obj.items())
    elif isinstance(obj, (list, tuple)):
        return map(pretty_convert, obj)
    return obj

def walk_obj(buf, obj, indent_num=2, depth=0, eol='', wrap_len=60, wrap_total=100, lf='\n'):
    indent = ' '*indent_num*depth
    buf.write(u'{0}'.format(eol if not eol else eol+indent) )
    if isinstance(obj, (basestring, int, float, complex)):
        buf.write(u'{0}'.format(obj).replace(lf, '{0}{1} '.format(lf, indent)) )
    elif isinstance(obj, ( dict )):
        eol, eol_org = lf, eol
        buf.write('{ ')
        for key in obj.keys():
            buf.write(u'{0}{1}{2}{3} '.format(eol, indent, ' '*indent_num, key))
            walk_obj(buf=buf, obj=obj[key], indent_num=indent_num, depth=depth+1, eol='', wrap_len=wrap_len, lf=lf )
            buf.write(',')
        buf.write('{0}}}'.format(eol if not eol else eol+indent) )
        eol=eol_org
    elif isinstance(obj, (list, tuple)):
        eol_org, indent_org = eol, indent
        for item in obj:
            if isinstance(item, (list, dict)):
                eol = lf
                break
            else:
                eol = ''
                indent = ''
                continue
        if max(map(len,obj)) > wrap_len or sum(map(len,obj)) > wrap_total:
            eol = lf
        buf.write('[ ')
        for item in obj:
            walk_obj(buf=buf, obj=item, indent_num=indent_num, depth=depth+1, eol=eol, wrap_len=wrap_len, lf=lf )
            buf.write(', ')
        buf.write('{0}]'.format(eol if not eol else eol+indent_org) )
        eol, indent = eol_org, indent_org
    if depth == 0:
        buf.write(lf)

def pretty_print(obj, indent=2, b=None):
    if not b:
        import sys, codecs
        b = codecs.getwriter('utf_8')(sys.stdout)
    walk_obj(b, pretty_convert(obj), indent_num=indent)


if __name__ == '__main__':
  import sys
  import codecs
  import yaml
  import pprint

  if len(sys.argv) < 2:
     print 'yaml file name is needed'
     sys.exit(2)
  indent = 2 if len(sys.argv) <3 else int(sys.argv[2])

  sys.stdout = codecs.getwriter('utf_8')(sys.stdout)

  file = codecs.open(sys.argv[1], 'r', 'utf_8').read()

  sys.stdout.write('# file \n\n```\n')
  sys.stdout.write(file)

  contents = yaml.load(file)

  sys.stdout.write('\n```\n\n# pprint.PrettyPrinter \n\n```\n\n')
  pp=pprint.PrettyPrinter(indent=indent)
  pp.pprint(contents)

  sys.stdout.write('\n```\n\n# pretty_print \n\n```\n')
  pretty_print(contents, indent, sys.stdout)
  sys.stdout.write('\n')

zz.yml

dictionary:
  - { test: this is test, test2: 1, test3: true }
  - [1,2,3,4,5]
  - - '日本語であそぼ'
    - [ 6,7,8,9,'テスト']
  - name: test
    名前: test2
  - test2
  - test3
  - true
  - 65535
  - 3.14
  - - list1
    - list2
    - list3
list:
  -  name: this is dictionary,
     bool: false
     int: 15000
     float: 15.6
     list: [ 0,1,2,3,4 ]
  - "this is test"
  - |
    this is test
    this is Test
    this is tesT