examples/yaml-hl/yaml_hl.py

   1 #!/usr/bin/python
   2
   3 import yaml, codecs, sys, os.path, optparse
   4
   5 class 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
  32 yaml.add_path_resolver(u'tag:yaml.org,2002:python/object:__main__.Style',
  33         [None], dict)
  34 yaml.add_path_resolver(u'tag:yaml.org,2002:pairs',
  35         [None, u'replaces'], list)
  36
  37 class YAMLHighlight:
  38
  39     def __init__(self, options):
  40         config = yaml.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
 100 if __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