examples/yaml-hl/yaml_hl.py

   1 #!/usr/bin/python
   2
   3 import yaml, codecs, sys, 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         if self.style.header:
  53             self.output.write(self.style.header)
  54         input = self.input.read()
  55         if input.startswith(codecs.BOM_UTF16_LE):
  56             input = unicode(input, 'utf-16-le')
  57         elif input.startswith(codecs.BOM_UTF16_BE):
  58             input = unicode(input, 'utf-16-be')
  59         else:
  60             input = unicode(input, 'utf-8')
  61         substitutions = self.style.substitutions
  62         tokens = yaml.scan(input)
  63         events = yaml.parse(input)
  64         markers = []
  65         number = 0
  66         for token in tokens:
  67             number += 1
  68             if token.start_mark.index != token.end_mark.index:
  69                 cls = token.__class__
  70                 if (cls, -1) in substitutions:
  71                     markers.append([token.start_mark.index, +2, number, substitutions[cls, -1]])
  72                 if (cls, +1) in substitutions:
  73                     markers.append([token.end_mark.index, -2, number, substitutions[cls, +1]])
  74         number = 0
  75         for event in events:
  76             number += 1
  77             cls = event.__class__
  78             if (cls, -1) in substitutions:
  79                 markers.append([event.start_mark.index, +1, number, substitutions[cls, -1]])
  80             if (cls, +1) in substitutions:
  81                 markers.append([event.end_mark.index, -1, number, substitutions[cls, +1]])
  82         markers.sort()
  83         markers.reverse()
  84         chunks = []
  85         position = len(input)
  86         for index, weight1, weight2, substitution in markers:
  87             if index < position:
  88                 chunk = input[index:position]
  89                 for substring, replacement in self.style.replaces:
  90                     chunk = chunk.replace(substring, replacement)
  91                 chunks.append(chunk)
  92                 position = index
  93             chunks.append(substitution)
  94         chunks.reverse()
  95         result = u''.join(chunks)
  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', default='yaml_hl.cfg',
 105             help="set an alternative configuration file", metavar='CONFIG')
 106     parser.add_option('-i', '--input', dest='input', default=None,
 107             help="set the input file (default: stdin)", metavar='FILE')
 108     parser.add_option('-o', '--output', dest='output', default=None,
 109             help="set the output file (default: stdout)", metavar='FILE')
 110     (options, args) = parser.parse_args()
 111     hl = YAMLHighlight(options)
 112     hl.highlight()
 113