1#!/usr/bin/python 2 3import yaml, codecs, sys, os.path, optparse 4 5class Style: 6 7 def __init__(self, header=None, footer=None, 8 tokens=None, events=None, replaces=None): 9 self.header = header 10 self.footer = footer 11 self.replaces = replaces 12 self.substitutions = {} 13 for domain, Class in [(tokens, 'Token'), (events, 'Event')]: 14 if not domain: 15 continue 16 for key in domain: 17 name = ''.join([part.capitalize() for part in key.split('-')]) 18 cls = getattr(yaml, '%s%s' % (name, Class)) 19 value = domain[key] 20 if not value: 21 continue 22 start = value.get('start') 23 end = value.get('end') 24 if start: 25 self.substitutions[cls, -1] = start 26 if end: 27 self.substitutions[cls, +1] = end 28 29 def __setstate__(self, state): 30 self.__init__(**state) 31 32yaml.add_path_resolver(u'tag:yaml.org,2002:python/object:__main__.Style', 33 [None], dict) 34yaml.add_path_resolver(u'tag:yaml.org,2002:pairs', 35 [None, u'replaces'], list) 36 37class YAMLHighlight: 38 39 def __init__(self, options): 40 config = yaml.full_load(file(options.config, 'rb').read()) 41 self.style = config[options.style] 42 if options.input: 43 self.input = file(options.input, 'rb') 44 else: 45 self.input = sys.stdin 46 if options.output: 47 self.output = file(options.output, 'wb') 48 else: 49 self.output = sys.stdout 50 51 def highlight(self): 52 input = self.input.read() 53 if input.startswith(codecs.BOM_UTF16_LE): 54 input = unicode(input, 'utf-16-le') 55 elif input.startswith(codecs.BOM_UTF16_BE): 56 input = unicode(input, 'utf-16-be') 57 else: 58 input = unicode(input, 'utf-8') 59 substitutions = self.style.substitutions 60 tokens = yaml.scan(input) 61 events = yaml.parse(input) 62 markers = [] 63 number = 0 64 for token in tokens: 65 number += 1 66 if token.start_mark.index != token.end_mark.index: 67 cls = token.__class__ 68 if (cls, -1) in substitutions: 69 markers.append([token.start_mark.index, +2, number, substitutions[cls, -1]]) 70 if (cls, +1) in substitutions: 71 markers.append([token.end_mark.index, -2, number, substitutions[cls, +1]]) 72 number = 0 73 for event in events: 74 number += 1 75 cls = event.__class__ 76 if (cls, -1) in substitutions: 77 markers.append([event.start_mark.index, +1, number, substitutions[cls, -1]]) 78 if (cls, +1) in substitutions: 79 markers.append([event.end_mark.index, -1, number, substitutions[cls, +1]]) 80 markers.sort() 81 markers.reverse() 82 chunks = [] 83 position = len(input) 84 for index, weight1, weight2, substitution in markers: 85 if index < position: 86 chunk = input[index:position] 87 for substring, replacement in self.style.replaces: 88 chunk = chunk.replace(substring, replacement) 89 chunks.append(chunk) 90 position = index 91 chunks.append(substitution) 92 chunks.reverse() 93 result = u''.join(chunks) 94 if self.style.header: 95 self.output.write(self.style.header) 96 self.output.write(result.encode('utf-8')) 97 if self.style.footer: 98 self.output.write(self.style.footer) 99 100if __name__ == '__main__': 101 parser = optparse.OptionParser() 102 parser.add_option('-s', '--style', dest='style', default='ascii', 103 help="specify the highlighting style", metavar='STYLE') 104 parser.add_option('-c', '--config', dest='config', 105 default=os.path.join(os.path.dirname(sys.argv[0]), 'yaml_hl.cfg'), 106 help="set an alternative configuration file", metavar='CONFIG') 107 parser.add_option('-i', '--input', dest='input', default=None, 108 help="set the input file (default: stdin)", metavar='FILE') 109 parser.add_option('-o', '--output', dest='output', default=None, 110 help="set the output file (default: stdout)", metavar='FILE') 111 (options, args) = parser.parse_args() 112 hl = YAMLHighlight(options) 113 hl.highlight() 114 115