# I/O

import timeouts

##### 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)

##### 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 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)