Project

General

Profile

1 1630 aaronmk
# I/O
2
3 1673 aaronmk
import timeouts
4 1635 aaronmk
5
class WrapStream:
6
    '''Forwards close() to the underlying stream'''
7
    def __init__(self, stream): self.stream = stream
8
9
    def close(self): self.stream.close()
10
11
class StreamIter(WrapStream):
12 1630 aaronmk
    '''Iterates over the lines in a stream.
13
    Unlike stream.__iter__(), doesn't wait for EOF to start returning lines.
14
    Offers curr() method.
15
    '''
16
    def __init__(self, stream):
17 1635 aaronmk
        WrapStream.__init__(self, stream)
18 1630 aaronmk
        self.line = None
19
20
    def __iter__(self): return self
21
22
    def curr(self):
23
        if self.line == None: self.line = self.stream.readline()
24
        if self.line == '': raise StopIteration
25
        return self.line
26
27
    def next(self):
28
        line = self.curr()
29
        self.line = None
30
        return line
31 1635 aaronmk
32
class TimeoutInputStream(WrapStream):
33 1673 aaronmk
    '''@param timeout sec'''
34
    def __init__(self, stream, timeout):
35 1635 aaronmk
        WrapStream.__init__(self, stream)
36 1673 aaronmk
        self.timeout = timeout
37 1630 aaronmk
38 1635 aaronmk
    def readline(self):
39 1673 aaronmk
        return timeouts.run(lambda: self.stream.readline(), self.timeout)
40 1630 aaronmk
41
class TracedOutputStream:
42
    '''Wraps an output stream, running a trace function on each string written
43
    '''
44
    def __init__(self, trace, stream):
45
        self.trace = trace
46
        self.stream = stream
47
48
    def write(self, str_):
49
        self.trace(str_)
50
        return self.stream.write(str_)
51
52
class LineCountOutputStream(TracedOutputStream):
53
    '''Wraps an output stream, making the current line number available.
54
    Lines start counting from 1.'''
55
    def __init__(self, stream):
56
        self.line_num = 1
57
        def trace(str_): self.line_num += str_.count('\n')
58
        TracedOutputStream.__init__(self, trace, stream)