Splitting YAML "directives" from other events, and passing the others down to yaml.load

pyyaml-split.py

__all__ = [ 'EventBasedLoader', 'EventBasedSafeLoader' ]

from yaml.composer import Composer
from yaml.constructor import Constructor, SafeConstructor
from yaml.resolver import Resolver
import yaml

# Replacement for Scanner+Parser that replays events from a stream.  These
# four methods seem to be the entire interface that Composer expects.
class EventBasedParser(object):
    def __init__(self, events):
        self.__events = events
        self.__current_event = None

    def check_event(self, *choices):
        # Check the type of the next event.
        if self.__current_event is None:
            try:
                self.__current_event = next(self.__events)
            except StopIteration:
                pass
        if self.__current_event is not None:
            if not choices:
                return True
            for choice in choices:
                if isinstance(self.__current_event, choice):
                    return True
        return False

    def dispose(self):
        pass

    def peek_event(self):
        # Get the next event.
        if self.__current_event is None:
            try:
                self.__current_event = next(self.__events)
            except StopIteration:
                pass
        return self.__current_event

    def get_event(self):
        # Get the next event and proceed further.
        if self.__current_event is None:
            try:
                self.__current_event = next(self.__events)
            except StopIteration:
                pass
        value = self.__current_event
        self.__current_event = None
        return value

class EventBasedLoader(EventBasedParser, Composer, Constructor, Resolver):
    """Same as the usual PyYAML Loader, but replaces the text-based Parser and
       Scanner with one that replays from an event stream."""
    def __init__(self, events):
        EventBasedParser.__init__(self, events)
        Composer.__init__(self)
        Constructor.__init__(self)
        Resolver.__init__(self)

class EventBasedSafeLoader(EventBasedParser, Composer, SafeConstructor, Resolver):
    """Same as the usual PyYAML SafeLoader, but replaces the text-based Parser and
       Scanner with an event stream."""
    def __init__(self, events):
        self.__events = events
        self.__current_event = None
        Composer.__init__(self)
        SafeConstructor.__init__(self)
        Resolver.__init__(self)


if __name__ == '__main__':
    def split_and_yield(events):
        for e in events:
            if isinstance(e, yaml.ScalarEvent) and not (e.tag is None) and e.tag.startswith('!mux-'):
                print e
            else:
                yield e

    doc = """
- without: 'tag'
  a: 'b'
  c: 'd'
- !mux-with 'tag'
"""

    # text -> events
    events = yaml.parse(doc)

    # events -> process tagged items, pass through the rest
    data_events = split_and_yield(events)

    # non-tagged items -> Python object
    data = yaml.load(data_events, EventBasedLoader)
    print(data)