# 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.readline()

        if self.line == '': raise StopIteration

        return self.line

    def next(self):

        line = self.curr()

        self.line = None

        return line

    # Define as a separate method so it can be overridden

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

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(StreamIter):

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

    def __init__(self, filter_, stream):

        StreamIter.__init__(self, stream)

        self.filter = filter_

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

    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_):

            if str_ != '': self.line_num += 1 # EOF is not a line

        TracedStream.__init__(self, trace, stream)

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

class ProgressInputStream(TracedStream):

    '''Wraps an input stream, reporting the # lines read every n lines and after

    the last line is read.

    @param log The output stream for progress messages

    '''

    def __init__(self, stream, log, msg='Read %d line(s)', n=100):

        self.eof = False

        def trace(str_):

            msg_ = msg # may be modified, so can't have same name as outer var

            if str_ == None: # closed

                if self.eof: return # not closed prematurely

                str_ = '' # closed prematurely, so signal EOF

                msg_ += ' (not all input read)'

            self.eof = eof = str_ == ''

            if eof: line_ending = '\n'

            else: line_ending = '\r'

            line_ct = self.stream.line_num - 1 # make it 0-based

            if eof or line_ct % n == 0: log.write((msg_ % line_ct)+line_ending)

        self.trace = trace

        TracedStream.__init__(self, trace, LineCountInputStream(stream))

    def close(self):

        self.trace(None) # signal closed

        TracedStream.close(self)

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)