# I/O

import timeouts

##### Exceptions

class InputStreamsOnlyException(Exception):

    def __init__(self): Exception.__init__(self, 'Only input streams allowed')

##### Wrapped streams

class WrapStream:

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

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

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

##### Line iteration

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

def copy(from_, to):

    for line in StreamIter(from_): to.write(line)

def consume(in_):

    for line in StreamIter(in_): pass

##### Timeouts

class TimeoutInputStream(WrapStream):

    '''@param timeout sec'''

    def __init__(self, stream, timeout):

        WrapStream.__init__(self, stream)

        self.timeout = timeout

    def readline(self):

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

    def write(self, str_): raise InputStreamsOnlyException()

##### Filtered/traced streams

class FilterStream(WrapStream):

    '''Wraps a stream, filtering each string read or written'''

    def __init__(self, filter_, stream):

        WrapStream.__init__(self, stream)

        self.filter = filter_

    def readline(self): return self.filter(self.stream.readline())

    def read(self, n): return self.readline() # forward all reads to readline()

    def write(self, str_): return self.stream.write(self.filter(str_))

class TracedStream(FilterStream):

    '''Wraps a stream, running a trace function on each string read or written

    '''

    def __init__(self, trace, stream):

        def filter_(str_):

            trace(str_)

            return str_

        FilterStream.__init__(self, filter_, stream)

class LineCountStream(TracedStream):

    '''Wraps a unidirectional 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')

        TracedStream.__init__(self, trace, stream)

class LineCountInputStream(TracedStream):

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

    Lines start counting from 1.

    This is faster than LineCountStream for input streams, because all reading

    is done through readline() so newlines don't need to be explicitly counted.

    '''

    def __init__(self, stream):

        self.line_num = 1

        def trace(str_): self.line_num += 1

        TracedStream.__init__(self, trace, stream)

    def write(self, str_): raise InputStreamsOnlyException()

class CaptureStream(TracedStream):

    '''Wraps a stream, capturing matching text.

    Matches can be retrieved from self.matches.'''

    def __init__(self, stream, start_str, end_str):

        self.recording = False

        self.matches = []

        def trace(str_):

            start_idx = 0

            if not self.recording:

                start_idx = str_.find(start_str)

                if start_idx >= 0:

                    self.recording = True

                    self.matches.append('')

            if self.recording:

                end_idx = str_.find(end_str)

                if end_idx >= 0:

                    self.recording = False

                    end_idx += len(end_str)

                else: end_idx = len(str_)

                self.matches[-1] += str_[start_idx:end_idx]

        TracedStream.__init__(self, trace, stream)