Project

General

Profile

1
# I/O
2

    
3
import timeouts
4

    
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
    '''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
        WrapStream.__init__(self, stream)
18
        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

    
32
class TimeoutInputStream(WrapStream):
33
    '''@param timeout sec'''
34
    def __init__(self, stream, timeout):
35
        WrapStream.__init__(self, stream)
36
        self.timeout = timeout
37
    
38
    def readline(self):
39
        return timeouts.run(lambda: self.stream.readline(), self.timeout)
40

    
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)
(13-13/22)