Project

General

Profile

1
# I/O
2

    
3
import operator
4
import timeouts
5

    
6
##### Exceptions
7

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

    
11
##### Wrapped streams
12

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

    
19
##### Line iteration
20

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

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

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

    
51
def read_all(stream): return reduce(operator.add, StreamIter(stream), '')
52

    
53
##### Timeouts
54

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

    
66
##### Filtered/traced streams
67

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

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

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

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

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

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