Project

General

Profile

1
# I/O
2

    
3
import timeouts
4

    
5
##### Exceptions
6

    
7
class InputStreamsOnlyException(Exception):
8
    def __init__(self): Exception.__init__(self, 'Only input streams allowed')
9

    
10
##### Wrapped streams
11

    
12
class WrapStream:
13
    '''Forwards close() to the underlying stream'''
14
    def __init__(self, stream): self.stream = stream
15
    
16
    def close(self): self.stream.close()
17

    
18
##### Line iteration
19

    
20
class StreamIter(WrapStream):
21
    '''Iterates over the lines in a stream.
22
    Unlike stream.__iter__(), doesn't wait for EOF to start returning lines.
23
    Offers curr() method.
24
    '''
25
    def __init__(self, stream):
26
        WrapStream.__init__(self, stream)
27
        self.line = None
28
    
29
    def __iter__(self): return self
30
    
31
    def curr(self):
32
        if self.line == None: self.line = self.readline()
33
        if self.line == '': raise StopIteration
34
        return self.line
35
    
36
    def next(self):
37
        line = self.curr()
38
        self.line = None
39
        return line
40
    
41
    # Define as a separate method so it can be overridden
42
    def readline(self): return self.stream.readline()
43

    
44
def copy(from_, to):
45
    for line in StreamIter(from_): to.write(line)
46

    
47
def consume(in_):
48
    for line in StreamIter(in_): pass
49

    
50
##### Timeouts
51

    
52
class TimeoutInputStream(WrapStream):
53
    '''@param timeout sec'''
54
    def __init__(self, stream, timeout):
55
        WrapStream.__init__(self, stream)
56
        self.timeout = timeout
57
    
58
    def readline(self):
59
        return timeouts.run(lambda: self.stream.readline(), self.timeout)
60
    
61
    def write(self, str_): raise InputStreamsOnlyException()
62

    
63
##### Filtered/traced streams
64

    
65
class FilterStream(StreamIter):
66
    '''Wraps a stream, filtering each string read or written'''
67
    def __init__(self, filter_, stream):
68
        StreamIter.__init__(self, stream)
69
        self.filter = filter_
70
    
71
    def readline(self): return self.filter(StreamIter.readline(self))
72
    
73
    def read(self, n): return self.readline() # forward all reads to readline()
74
    
75
    def write(self, str_): return self.stream.write(self.filter(str_))
76

    
77
class TracedStream(FilterStream):
78
    '''Wraps a stream, running a trace function on each string read or written
79
    '''
80
    def __init__(self, trace, stream):
81
        def filter_(str_):
82
            trace(str_)
83
            return str_
84
        FilterStream.__init__(self, filter_, stream)
85

    
86
class LineCountStream(TracedStream):
87
    '''Wraps a unidirectional stream, making the current line number available.
88
    Lines start counting from 1.'''
89
    def __init__(self, stream):
90
        self.line_num = 1
91
        def trace(str_): self.line_num += str_.count('\n')
92
        TracedStream.__init__(self, trace, stream)
93

    
94
class LineCountInputStream(TracedStream):
95
    '''Wraps an input stream, making the current line number available.
96
    Lines start counting from 1.
97
    This is faster than LineCountStream for input streams, because all reading
98
    is done through readline() so newlines don't need to be explicitly counted.
99
    '''
100
    def __init__(self, stream):
101
        self.line_num = 1
102
        def trace(str_): self.line_num += 1
103
        TracedStream.__init__(self, trace, stream)
104
    
105
    def write(self, str_): raise InputStreamsOnlyException()
106

    
107
class ProgressInputStream(TracedStream):
108
    '''Wraps an input stream, reporting the # lines read every n lines and after
109
    the last line is read.
110
    @param log The output stream for progress messages
111
    '''
112
    def __init__(self, stream, log, msg='Read %d line(s)', n=100):
113
        self.eof = False
114
        
115
        def trace(str_):
116
            msg_ = msg # may be modified, so can't have same name as outer var
117
            if str_ == None: # closed
118
                if self.eof: return # not closed prematurely
119
                str_ = '' # closed prematurely, so signal EOF
120
                msg_ += ' (not all input read)'
121
            
122
            self.eof = eof = str_ == ''
123
            if eof: line_ending = '\n'
124
            else: line_ending = '\r'
125
            
126
            line_ct = self.stream.line_num - 1 # make it 0-based
127
            if eof or line_ct % n == 0: log.write((msg_ % line_ct)+line_ending)
128
        self.trace = trace
129
        
130
        TracedStream.__init__(self, trace, LineCountInputStream(stream))
131
    
132
    def close(self):
133
        self.trace(None) # signal closed
134
        TracedStream.close(self)
135

    
136
class CaptureStream(TracedStream):
137
    '''Wraps a stream, capturing matching text.
138
    Matches can be retrieved from self.matches.'''
139
    def __init__(self, stream, start_str, end_str):
140
        self.recording = False
141
        self.matches = []
142
        
143
        def trace(str_):
144
            start_idx = 0
145
            if not self.recording:
146
                start_idx = str_.find(start_str)
147
                if start_idx >= 0:
148
                    self.recording = True
149
                    self.matches.append('')
150
            if self.recording:
151
                end_idx = str_.find(end_str)
152
                if end_idx >= 0:
153
                    self.recording = False
154
                    end_idx += len(end_str)
155
                else: end_idx = len(str_)
156
                
157
                self.matches[-1] += str_[start_idx:end_idx]
158
        
159
        TracedStream.__init__(self, trace, stream)
(26-26/36)