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 2672 aaronmk
        def trace(str_):
103
            if str_ != '': self.line_num += 1 # EOF is not a line
104 1928 aaronmk
        TracedStream.__init__(self, trace, stream)
105
106
    def write(self, str_): raise InputStreamsOnlyException()
107
108 1939 aaronmk
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 1682 aaronmk
class CaptureStream(TracedStream):
138
    '''Wraps a stream, capturing matching text.
139 1706 aaronmk
    Matches can be retrieved from self.matches.'''
140 1682 aaronmk
    def __init__(self, stream, start_str, end_str):
141
        self.recording = False
142 1706 aaronmk
        self.matches = []
143 1682 aaronmk
144
        def trace(str_):
145 1707 aaronmk
            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 1682 aaronmk
160
        TracedStream.__init__(self, trace, stream)