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_):
103
            if str_ != '': self.line_num += 1 # EOF is not a line
104
        TracedStream.__init__(self, trace, stream)
105
    
106
    def write(self, str_): raise InputStreamsOnlyException()
107

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

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