# I/O

import timeout

class WrapStream:

    '''Forwards close() to the underlying stream'''

    def __init__(self, stream): self.stream = stream

    def close(self): self.stream.close()

class StreamIter(WrapStream):

    '''Iterates over the lines in a stream.

    Unlike stream.__iter__(), doesn't wait for EOF to start returning lines.

    Offers curr() method.

    '''

    def __init__(self, stream):

        WrapStream.__init__(self, stream)

        self.line = None

    def __iter__(self): return self

    def curr(self):

        if self.line == None: self.line = self.stream.readline()

        if self.line == '': raise StopIteration

        return self.line

    def next(self):

        line = self.curr()

        self.line = None

        return line

class TimeoutInputStream(WrapStream):

    '''@param timeout_ sec'''

    def __init__(self, stream, timeout_):

        WrapStream.__init__(self, stream)

        self.timeout = timeout_

    def readline(self):

        return timeout.run(lambda: self.stream.readline(), self.timeout)

class TracedOutputStream:

    '''Wraps an output stream, running a trace function on each string written

    '''

    def __init__(self, trace, stream):

        self.trace = trace

        self.stream = stream

    def write(self, str_):

        self.trace(str_)

        return self.stream.write(str_)

class LineCountOutputStream(TracedOutputStream):

    '''Wraps an output stream, making the current line number available.

    Lines start counting from 1.'''

    def __init__(self, stream):

        self.line_num = 1

        def trace(str_): self.line_num += str_.count('\n')

        TracedOutputStream.__init__(self, trace, stream)