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