Source

pyyaml / examples / yaml-hl / yaml_hl.py

Diff from to

File examples/yaml-hl/yaml_hl.py

 
 import yaml, codecs, sys, optparse
 
-#yaml.add_resolver(u'!Config', [])
-#yaml.add_resolver(u'!TokensConfig', [u'tokens'])
-#yaml.add_resolver(u'!EventsConfig', [u'events'])
-#yaml.add_resolver(u'!StartEndConfig', [u'tokens', None])
-#yaml.add_resolver(u'!StartEndConfig', [u'events', None])
+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, config):
-        parameters = yaml.load(config)
-        self.replaces = parameters['replaces']
-        self.substitutions = {}
-        for domain, items in [('Token', parameters['tokens']),
-                    ('Event', parameters['events'])]:
-            for code in items:
-                name = ''.join([part.capitalize() for part in code.split('-')]+[domain])
-                cls = getattr(yaml, name)
-                value = items[code]
-                if value:
-                    if 'start' in value:
-                        self.substitutions[cls, -1] = value['start']
-                    if 'end' in value:
-                        self.substitutions[cls, +1] = value['end']
+    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):
-        if isinstance(input, str):
-            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')
+    def highlight(self):
+        if self.style.header:
+            self.output.write(self.style.header)
+        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 += 1
             if token.start_mark.index != token.end_mark.index:
                 cls = token.__class__
-                if (cls, -1) in self.substitutions:
-                    markers.append([token.start_mark.index, +2, number, self.substitutions[cls, -1]])
-                if (cls, +1) in self.substitutions:
-                    markers.append([token.end_mark.index, -2, number, self.substitutions[cls, +1]])
+                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 self.substitutions:
-                markers.append([event.start_mark.index, +1, number, self.substitutions[cls, -1]])
-            if (cls, +1) in self.substitutions:
-                markers.append([event.end_mark.index, -1, number, self.substitutions[cls, +1]])
+            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 = []
         for index, weight1, weight2, substitution in markers:
             if index < position:
                 chunk = input[index:position]
-                for substring, replacement in self.replaces:
+                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)
-        return result.encode('utf-8')
+        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('-c', '--config', dest='config', default='yaml_hl_ascii.cfg', metavar='CONFIG')
+    parser.add_option('-s', '--style', dest='style', default='ascii',
+            help="specify the highlighting style", metavar='STYLE')
+    parser.add_option('-c', '--config', dest='config', default='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(file(options.config))
-    sys.stdout.write(hl.highlight(sys.stdin.read()))
+    hl = YAMLHighlight(options)
+    hl.highlight()