Project

General

Profile

1 1630 aaronmk
# I/O
2
3 1673 aaronmk
import timeouts
4 1635 aaronmk
5 1928 aaronmk
##### Exceptions
6
7
class InputStreamsOnlyException(Exception):
8
    def __init__(self): Exception.__init__(self, 'Only input streams allowed')
9
10 1759 aaronmk
##### Wrapped streams
11
12 1635 aaronmk
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 1759 aaronmk
##### Line iteration
19
20 1635 aaronmk
class StreamIter(WrapStream):
21 1630 aaronmk
    '''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 1635 aaronmk
        WrapStream.__init__(self, stream)
27 1630 aaronmk
        self.line = None
28
29
    def __iter__(self): return self
30
31
    def curr(self):
32 1964 aaronmk
        if self.line == None: self.line = self.readline()
33 1630 aaronmk
        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 1964 aaronmk
41
    # Define as a separate method so it can be overridden
42
    def readline(self): return self.stream.readline()
43 1635 aaronmk
44 1686 aaronmk
def copy(from_, to):
45
    for line in StreamIter(from_): to.write(line)
46
47 1759 aaronmk
def consume(in_):
48
    for line in StreamIter(in_): pass
49
50
##### Timeouts
51
52 1635 aaronmk
class TimeoutInputStream(WrapStream):
53 1673 aaronmk
    '''@param timeout sec'''
54
    def __init__(self, stream, timeout):
55 1635 aaronmk
        WrapStream.__init__(self, stream)
56 1673 aaronmk
        self.timeout = timeout
57 1630 aaronmk
58 1635 aaronmk
    def readline(self):
59 1673 aaronmk
        return timeouts.run(lambda: self.stream.readline(), self.timeout)
60 1928 aaronmk
61
    def write(self, str_): raise InputStreamsOnlyException()
62 1630 aaronmk
63 1710 aaronmk
##### Filtered/traced streams
64 1686 aaronmk
65 1962 aaronmk
class FilterStream(StreamIter):
66 1710 aaronmk
    '''Wraps a stream, filtering each string read or written'''
67
    def __init__(self, filter_, stream):
68 1962 aaronmk
        StreamIter.__init__(self, stream)
69 1710 aaronmk
        self.filter = filter_
70 1630 aaronmk
71 1964 aaronmk
    def readline(self): return self.filter(StreamIter.readline(self))
72 1682 aaronmk
73 1762 aaronmk
    def read(self, n): return self.readline() # forward all reads to readline()
74
75 1710 aaronmk
    def write(self, str_): return self.stream.write(self.filter(str_))
76 1630 aaronmk
77 1710 aaronmk
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 1684 aaronmk
class LineCountStream(TracedStream):
87
    '''Wraps a unidirectional stream, making the current line number available.
88 1630 aaronmk
    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 1682 aaronmk
        TracedStream.__init__(self, trace, stream)
93
94 1928 aaronmk
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 1939 aaronmk
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 1682 aaronmk
class CaptureStream(TracedStream):
137
    '''Wraps a stream, capturing matching text.
138 1706 aaronmk
    Matches can be retrieved from self.matches.'''
139 1682 aaronmk
    def __init__(self, stream, start_str, end_str):
140
        self.recording = False
141 1706 aaronmk
        self.matches = []
142 1682 aaronmk
143
        def trace(str_):
144 1707 aaronmk
            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 1682 aaronmk
159
        TracedStream.__init__(self, trace, stream)