tenfourfox/python/pyyaml/examples/yaml-highlight/yaml_hl.py

#!/usr/bin/python

import yaml, codecs, sys, os.path, optparse

class Style:

    def __init__(self, header=None, footer=None,
            tokens=None, events=None, replaces=None):
        self.header = header
        self.footer = footer
        self.replaces = replaces
        self.substitutions = {}
        for domain, Class in [(tokens, 'Token'), (events, 'Event')]:
            if not domain:
                continue
            for key in domain:
                name = ''.join([part.capitalize() for part in key.split('-')])
                cls = getattr(yaml, '%s%s' % (name, Class))
                value = domain[key]
                if not value:
                    continue
                start = value.get('start')
                end = value.get('end')
                if start:
                    self.substitutions[cls, -1] = start
                if end:
                    self.substitutions[cls, +1] = end

    def __setstate__(self, state):
        self.__init__(**state)

yaml.add_path_resolver(u'tag:yaml.org,2002:python/object:__main__.Style',
        [None], dict)
yaml.add_path_resolver(u'tag:yaml.org,2002:pairs',
        [None, u'replaces'], list)

class YAMLHighlight:

    def __init__(self, options):
        config = yaml.load(file(options.config, 'rb').read())
        self.style = config[options.style]
        if options.input:
            self.input = file(options.input, 'rb')
        else:
            self.input = sys.stdin
        if options.output:
            self.output = file(options.output, 'wb')
        else:
            self.output = sys.stdout

    def highlight(self):
        input = self.input.read()
        if input.startswith(codecs.BOM_UTF16_LE):
            input = unicode(input, 'utf-16-le')
        elif input.startswith(codecs.BOM_UTF16_BE):
            input = unicode(input, 'utf-16-be')
        else:
            input = unicode(input, 'utf-8')
        substitutions = self.style.substitutions
        tokens = yaml.scan(input)
        events = yaml.parse(input)
        markers = []
        number = 0
        for token in tokens:
            number += 1
            if token.start_mark.index != token.end_mark.index:
                cls = token.__class__
                if (cls, -1) in substitutions:
                    markers.append([token.start_mark.index, +2, number, substitutions[cls, -1]])
                if (cls, +1) in substitutions:
                    markers.append([token.end_mark.index, -2, number, substitutions[cls, +1]])
        number = 0
        for event in events:
            number += 1
            cls = event.__class__
            if (cls, -1) in substitutions:
                markers.append([event.start_mark.index, +1, number, substitutions[cls, -1]])
            if (cls, +1) in substitutions:
                markers.append([event.end_mark.index, -1, number, substitutions[cls, +1]])
        markers.sort()
        markers.reverse()
        chunks = []
        position = len(input)
        for index, weight1, weight2, substitution in markers:
            if index < position:
                chunk = input[index:position]
                for substring, replacement in self.style.replaces:
                    chunk = chunk.replace(substring, replacement)
                chunks.append(chunk)
                position = index
            chunks.append(substitution)
        chunks.reverse()
        result = u''.join(chunks)
        if self.style.header:
            self.output.write(self.style.header)
        self.output.write(result.encode('utf-8'))
        if self.style.footer:
            self.output.write(self.style.footer)

if __name__ == '__main__':
    parser = optparse.OptionParser()
    parser.add_option('-s', '--style', dest='style', default='ascii',
            help="specify the highlighting style", metavar='STYLE')
    parser.add_option('-c', '--config', dest='config',
            default=os.path.join(os.path.dirname(sys.argv[0]), 'yaml_hl.cfg'),
            help="set an alternative configuration file", metavar='CONFIG')
    parser.add_option('-i', '--input', dest='input', default=None,
            help="set the input file (default: stdin)", metavar='FILE')
    parser.add_option('-o', '--output', dest='output', default=None,
            help="set the output file (default: stdout)", metavar='FILE')
    (options, args) = parser.parse_args()
    hl = YAMLHighlight(options)
    hl.highlight()