bundled/werkzeug/werkzeug/contrib/iterio.py @ b2bd9c232faa flask

# -*- coding: utf-8 -*-
r"""
    werkzeug.contrib.iterio
    ~~~~~~~~~~~~~~~~~~~~~~~

    This module implements a :class:`IterIO` that converts an iterator into
    a stream object and the other way round.  Converting streams into
    iterators requires the `greenlet`_ module.

    To convert an iterator into a stream all you have to do is to pass it
    directly to the :class:`IterIO` constructor.  In this example we pass it
    a newly created generator::

        def foo():
            yield "something\n"
            yield "otherthings"
        stream = IterIO(foo())
        print stream.read()         # read the whole iterator

    The other way round works a bit different because we have to ensure that
    the code execution doesn't take place yet.  An :class:`IterIO` call with a
    callable as first argument does two things.  The function itself is passed
    an :class:`IterIO` stream it can feed.  The object returned by the
    :class:`IterIO` constructor on the other hand is not an stream object but
    an iterator::

        def foo(stream):
            stream.write("some")
            stream.write("thing")
            stream.flush()
            stream.write("otherthing")
        iterator = IterIO(foo)
        print iterator.next()       # prints something
        print iterator.next()       # prints otherthing
        iterator.next()             # raises StopIteration

    .. _greenlet: http://codespeak.net/py/dist/greenlet.html

    :copyright: (c) 2010 by the Werkzeug Team, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
"""
try:
    from py.magic import greenlet
except:
    greenlet = None


class IterIO(object):
    """Instances of this object implement an interface compatible with the
    standard Python :class:`file` object.  Streams are either read-only or
    write-only depending on how the object is created.
    """

    def __new__(cls, obj):
        try:
            iterator = iter(obj)
        except TypeError:
            return IterI(obj)
        return IterO(iterator)

    def __iter__(self):
        return self

    def tell(self):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        return self.pos

    def isatty(self):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        return False

    def seek(self, pos, mode=0):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def truncate(self, size=None):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def write(self, s):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def writelines(self, list):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def read(self, n=-1):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def readlines(self, sizehint=0):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def readline(self, length=None):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def flush(self):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        raise IOError(9, 'Bad file descriptor')

    def next(self):
        if self.closed:
            raise StopIteration()
        line = self.readline()
        if not line:
            raise StopIteration()
        return line


class IterI(IterIO):
    """Convert an stream into an iterator."""

    def __new__(cls, func):
        if greenlet is None:
            raise RuntimeError('IterI requires greenlet support')
        stream = object.__new__(cls)
        stream.__init__(greenlet.getcurrent())

        def run():
            func(stream)
            stream.flush()

        g = greenlet(run, stream._parent)
        while 1:
            rv = g.switch()
            if not rv:
                return
            yield rv[0]

    def __init__(self, parent):
        self._parent = parent
        self._buffer = []
        self.closed = False
        self.pos = 0

    def close(self):
        if not self.closed:
            self.closed = True

    def write(self, s):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        self.pos += len(s)
        self._buffer.append(s)

    def writelines(self, list):
        self.write(''.join(list))

    def flush(self):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        data = ''.join(self._buffer)
        self._buffer = []
        self._parent.switch((data,))


class IterO(IterIO):
    """Iter output.  Wrap an iterator and give it a stream like interface."""

    def __new__(cls, gen):
        return object.__new__(cls)

    def __init__(self, gen):
        self._gen = gen
        self._buf = ''
        self.closed = False
        self.pos = 0

    def __iter__(self):
        return self

    def close(self):
        if not self.closed:
            self.closed = True
            if hasattr(self._gen, 'close'):
                self._gen.close()

    def seek(self, pos, mode=0):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        if mode == 1:
            pos += self.pos
        elif mode == 2:
            self.read()
            self.pos = min(self.pos, self.pos + pos)
            return
        elif mode != 0:
            raise IOError('Invalid argument')
        buf = []
        try:
            tmp_end_pos = len(self._buf)
            while pos > tmp_end_pos:
                item = self._gen.next()
                tmp_end_pos += len(item)
                buf.append(item)
        except StopIteration:
            pass
        if buf:
            self._buf += ''.join(buf)
        self.pos = max(0, pos)

    def read(self, n=-1):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        if n < 0:
            self._buf += ''.join(self._gen)
            result = self._buf[self.pos:]
            self.pos += len(result)
            return result
        new_pos = self.pos + n
        buf = []
        try:
            tmp_end_pos = len(self._buf)
            while new_pos > tmp_end_pos:
                item = self._gen.next()
                tmp_end_pos += len(item)
                buf.append(item)
        except StopIteration:
            pass
        if buf:
            self._buf += ''.join(buf)
        new_pos = max(0, new_pos)
        try:
            return self._buf[self.pos:new_pos]
        finally:
            self.pos = min(new_pos, len(self._buf))

    def readline(self, length=None):
        if self.closed:
            raise ValueError('I/O operation on closed file')
        nl_pos = self._buf.find('\n', self.pos)
        buf = []
        try:
            pos = self.pos
            while nl_pos < 0:
                item = self._gen.next()
                local_pos = item.find('\n')
                buf.append(item)
                if local_pos >= 0:
                    nl_pos = pos + local_pos
                    break
                pos += len(item)
        except StopIteration:
            pass
        if buf:
            self._buf += ''.join(buf)
        if nl_pos < 0:
            new_pos = len(self._buf)
        else:
            new_pos = nl_pos + 1
        if length is not None and self.pos + length < new_pos:
            new_pos = self.pos + length
        try:
            return self._buf[self.pos:new_pos]
        finally:
            self.pos = min(new_pos, len(self._buf))

    def readlines(self, sizehint=0):
        total = 0
        lines = []
        line = self.readline()
        while line:
            lines.append(line)
            total += len(line)
            if 0 < sizehint <= total:
                break
            line = self.readline()
        return lines