Project

General

Profile

1 1630 aaronmk
# I/O
2
3 1673 aaronmk
import timeouts
4 1635 aaronmk
5 1759 aaronmk
##### Wrapped streams
6
7 1635 aaronmk
class WrapStream:
8
    '''Forwards close() to the underlying stream'''
9
    def __init__(self, stream): self.stream = stream
10
11
    def close(self): self.stream.close()
12
13 1759 aaronmk
##### Line iteration
14
15 1635 aaronmk
class StreamIter(WrapStream):
16 1630 aaronmk
    '''Iterates over the lines in a stream.
17
    Unlike stream.__iter__(), doesn't wait for EOF to start returning lines.
18
    Offers curr() method.
19
    '''
20
    def __init__(self, stream):
21 1635 aaronmk
        WrapStream.__init__(self, stream)
22 1630 aaronmk
        self.line = None
23
24
    def __iter__(self): return self
25
26
    def curr(self):
27
        if self.line == None: self.line = self.stream.readline()
28
        if self.line == '': raise StopIteration
29
        return self.line
30
31
    def next(self):
32
        line = self.curr()
33
        self.line = None
34
        return line
35 1635 aaronmk
36 1686 aaronmk
def copy(from_, to):
37
    for line in StreamIter(from_): to.write(line)
38
39 1759 aaronmk
def consume(in_):
40
    for line in StreamIter(in_): pass
41
42
##### Timeouts
43
44 1635 aaronmk
class TimeoutInputStream(WrapStream):
45 1673 aaronmk
    '''@param timeout sec'''
46
    def __init__(self, stream, timeout):
47 1635 aaronmk
        WrapStream.__init__(self, stream)
48 1673 aaronmk
        self.timeout = timeout
49 1630 aaronmk
50 1635 aaronmk
    def readline(self):
51 1673 aaronmk
        return timeouts.run(lambda: self.stream.readline(), self.timeout)
52 1630 aaronmk
53 1710 aaronmk
##### Filtered/traced streams
54 1686 aaronmk
55 1710 aaronmk
class FilterStream(WrapStream):
56
    '''Wraps a stream, filtering each string read or written'''
57
    def __init__(self, filter_, stream):
58 1682 aaronmk
        WrapStream.__init__(self, stream)
59 1710 aaronmk
        self.filter = filter_
60 1630 aaronmk
61 1710 aaronmk
    def readline(self): return self.filter(self.stream.readline())
62 1682 aaronmk
63 1762 aaronmk
    def read(self, n): return self.readline() # forward all reads to readline()
64
65 1710 aaronmk
    def write(self, str_): return self.stream.write(self.filter(str_))
66 1630 aaronmk
67 1710 aaronmk
class TracedStream(FilterStream):
68
    '''Wraps a stream, running a trace function on each string read or written
69
    '''
70
    def __init__(self, trace, stream):
71
        def filter_(str_):
72
            trace(str_)
73
            return str_
74
        FilterStream.__init__(self, filter_, stream)
75
76 1684 aaronmk
class LineCountStream(TracedStream):
77
    '''Wraps a unidirectional stream, making the current line number available.
78 1630 aaronmk
    Lines start counting from 1.'''
79
    def __init__(self, stream):
80
        self.line_num = 1
81
        def trace(str_): self.line_num += str_.count('\n')
82 1682 aaronmk
        TracedStream.__init__(self, trace, stream)
83
84
class CaptureStream(TracedStream):
85
    '''Wraps a stream, capturing matching text.
86 1706 aaronmk
    Matches can be retrieved from self.matches.'''
87 1682 aaronmk
    def __init__(self, stream, start_str, end_str):
88
        self.recording = False
89 1706 aaronmk
        self.matches = []
90 1682 aaronmk
91
        def trace(str_):
92 1707 aaronmk
            start_idx = 0
93
            if not self.recording:
94
                start_idx = str_.find(start_str)
95
                if start_idx >= 0:
96
                    self.recording = True
97
                    self.matches.append('')
98
            if self.recording:
99
                end_idx = str_.find(end_str)
100
                if end_idx >= 0:
101
                    self.recording = False
102
                    end_idx += len(end_str)
103
                else: end_idx = len(str_)
104
105
                self.matches[-1] += str_[start_idx:end_idx]
106 1682 aaronmk
107
        TracedStream.__init__(self, trace, stream)