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
def copy(from_, to):
33
    for line in StreamIter(from_): to.write(line)
34

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

    
44
##### Filtered/traced streams
45

    
46
class FilterStream(WrapStream):
47
    '''Wraps a stream, filtering each string read or written'''
48
    def __init__(self, filter_, stream):
49
        WrapStream.__init__(self, stream)
50
        self.filter = filter_
51
    
52
    def readline(self): return self.filter(self.stream.readline())
53
    
54
    def write(self, str_): return self.stream.write(self.filter(str_))
55

    
56
class TracedStream(FilterStream):
57
    '''Wraps a stream, running a trace function on each string read or written
58
    '''
59
    def __init__(self, trace, stream):
60
        def filter_(str_):
61
            trace(str_)
62
            return str_
63
        FilterStream.__init__(self, filter_, stream)
64

    
65
class LineCountStream(TracedStream):
66
    '''Wraps a unidirectional stream, making the current line number available.
67
    Lines start counting from 1.'''
68
    def __init__(self, stream):
69
        self.line_num = 1
70
        def trace(str_): self.line_num += str_.count('\n')
71
        TracedStream.__init__(self, trace, stream)
72

    
73
class CaptureStream(TracedStream):
74
    '''Wraps a stream, capturing matching text.
75
    Matches can be retrieved from self.matches.'''
76
    def __init__(self, stream, start_str, end_str):
77
        self.recording = False
78
        self.matches = []
79
        
80
        def trace(str_):
81
            start_idx = 0
82
            if not self.recording:
83
                start_idx = str_.find(start_str)
84
                if start_idx >= 0:
85
                    self.recording = True
86
                    self.matches.append('')
87
            if self.recording:
88
                end_idx = str_.find(end_str)
89
                if end_idx >= 0:
90
                    self.recording = False
91
                    end_idx += len(end_str)
92
                else: end_idx = len(str_)
93
                
94
                self.matches[-1] += str_[start_idx:end_idx]
95
        
96
        TracedStream.__init__(self, trace, stream)
(15-15/25)