Project

General

Profile

1 1388 aaronmk
# CSV I/O
2
3
import csv
4 5431 aaronmk
import _csv
5 1388 aaronmk
import StringIO
6
7 8069 aaronmk
import exc
8 5593 aaronmk
import streams
9 1623 aaronmk
import strings
10 1444 aaronmk
import util
11
12 5426 aaronmk
delims = ',;\t|`'
13 5439 aaronmk
tab_padded_delims = ['\t|\t']
14 1623 aaronmk
tsv_delim = '\t'
15
escape = '\\'
16 1442 aaronmk
17 1623 aaronmk
ending_placeholder = r'\n'
18
19 5437 aaronmk
def is_tsv(dialect): return dialect.delimiter.startswith(tsv_delim)
20 1623 aaronmk
21 1621 aaronmk
def sniff(line):
22
    '''Automatically detects the dialect'''
23 5435 aaronmk
    line, ending = strings.extract_line_ending(line)
24 9961 aaronmk
    try: dialect = csv.Sniffer().sniff(line, delims)
25
    except _csv.Error, e:
26
        if exc.e_msg(e) == 'Could not determine delimiter': dialect = csv.excel
27
        else: raise
28 5435 aaronmk
29 5439 aaronmk
    if is_tsv(dialect):
30 8070 aaronmk
        dialect.quoting = csv.QUOTE_NONE
31 5439 aaronmk
        # Check multi-char delims using \t
32
        delim = strings.find_any(line, tab_padded_delims)
33
        if delim:
34
            dialect.delimiter = delim
35
            line_suffix = delim.rstrip('\t')
36
            if line.endswith(line_suffix): ending = line_suffix+ending
37 1621 aaronmk
    else: dialect.doublequote = True # Sniffer doesn't turn this on by default
38 5435 aaronmk
    dialect.lineterminator = ending
39
40 1621 aaronmk
    return dialect
41
42 8202 aaronmk
def has_unbalanced_quotes(str_): return str_.count('"') % 2 == 1 # odd # of "
43
44
def has_multiline_column(str_): return has_unbalanced_quotes(str_)
45
46 1923 aaronmk
def stream_info(stream, parse_header=False):
47 6589 aaronmk
    '''Automatically detects the dialect based on the header line.
48
    Uses the Excel dialect if the CSV file is empty.
49 1923 aaronmk
    @return NamedTuple {header_line, header, dialect}'''
50 1444 aaronmk
    info = util.NamedTuple()
51
    info.header_line = stream.readline()
52 8202 aaronmk
    if has_multiline_column(info.header_line): # 1st line not full header
53
        # assume it's a header-only csv with multiline columns
54
        info.header_line += ''.join(stream.readlines()) # use entire file
55 1923 aaronmk
    info.header = None
56
    if info.header_line != '':
57
        info.dialect = sniff(info.header_line)
58 6589 aaronmk
    else: info.dialect = csv.excel # line of '' indicates EOF = empty stream
59
60
    if parse_header:
61
        try: info.header = reader_class(info.dialect)(
62
            StringIO.StringIO(info.header_line), info.dialect).next()
63
        except StopIteration: info.header = []
64
65 1444 aaronmk
    return info
66
67 5170 aaronmk
tsv_encode_map = strings.json_encode_map[:]
68
tsv_encode_map.append(('\t', r'\t'))
69
tsv_decode_map = strings.flip_map(tsv_encode_map)
70 5146 aaronmk
71 1623 aaronmk
class TsvReader:
72
    '''Unlike csv.reader, for TSVs, interprets \ as escaping a line ending but
73 5145 aaronmk
    ignores it before everything else (e.g. \N for NULL).
74 5170 aaronmk
    Also expands tsv_encode_map escapes.
75 5145 aaronmk
    '''
76 1623 aaronmk
    def __init__(self, stream, dialect):
77
        assert is_tsv(dialect)
78
        self.stream = stream
79
        self.dialect = dialect
80
81
    def __iter__(self): return self
82
83
    def next(self):
84
        record = ''
85
        ending = None
86
        while True:
87
            line = self.stream.readline()
88
            if line == '': raise StopIteration
89
90 5438 aaronmk
            line = strings.remove_suffix(self.dialect.lineterminator, line)
91 5433 aaronmk
            contents = strings.remove_suffix(escape, line)
92 1623 aaronmk
            record += contents
93 5433 aaronmk
            if len(contents) == len(line): break # no line continuation
94 1623 aaronmk
            record += ending_placeholder
95 3055 aaronmk
96
        # Prevent "new-line character seen in unquoted field" errors
97
        record = record.replace('\r', ending_placeholder)
98
99 7211 aaronmk
        # Split line
100 8071 aaronmk
        if record == '': row = [] # csv.reader would interpret as EOF
101
        elif len(self.dialect.delimiter) > 1: # multi-char delims
102 7211 aaronmk
            row = record.split(self.dialect.delimiter)
103
        else: row = csv.reader(StringIO.StringIO(record), self.dialect).next()
104
105 5170 aaronmk
        return [strings.replace_all(tsv_decode_map, v) for v in row]
106 1623 aaronmk
107
def reader_class(dialect):
108
    if is_tsv(dialect): return TsvReader
109
    else: return csv.reader
110
111
def make_reader(stream, dialect): return reader_class(dialect)(stream, dialect)
112
113 1388 aaronmk
def reader_and_header(stream):
114
    '''Automatically detects the dialect based on the header line
115
    @return tuple (reader, header)'''
116 1923 aaronmk
    info = stream_info(stream, parse_header=True)
117 1958 aaronmk
    return (make_reader(stream, info.dialect), info.header)
118 1446 aaronmk
119 14577 aaronmk
def header(stream):
120
    '''fetches just the header line of a stream'''
121
    reader, header = reader_and_header(stream)
122
    return header
123
124 1446 aaronmk
##### csv modifications
125
126
# Note that these methods only work on *instances* of Dialect classes
127
csv.Dialect.__eq__ = lambda self, other: self.__dict__ == other.__dict__
128
csv.Dialect.__ne__ = lambda self, other: not (self == other)
129 2114 aaronmk
130 5430 aaronmk
__Dialect__validate_orig = csv.Dialect._validate
131
def __Dialect__validate(self):
132
        try: __Dialect__validate_orig(self)
133
        except _csv.Error, e:
134
            if str(e) == '"delimiter" must be an 1-character string': pass # OK
135
            else: raise
136
csv.Dialect._validate = __Dialect__validate
137
138 2114 aaronmk
##### Row filters
139
140
class Filter:
141
    '''Wraps a reader, filtering each row'''
142
    def __init__(self, filter_, reader):
143
        self.reader = reader
144
        self.filter = filter_
145
146
    def __iter__(self): return self
147
148
    def next(self): return self.filter(self.reader.next())
149 5574 aaronmk
150
    def close(self): pass # support using as a stream
151 2114 aaronmk
152
std_nulls = [r'\N']
153
empty_nulls = [''] + std_nulls
154
155
class NullFilter(Filter):
156
    '''Translates special string values to None'''
157
    def __init__(self, reader, nulls=std_nulls):
158
        map_ = dict.fromkeys(nulls, None)
159
        def filter_(row): return [map_.get(v, v) for v in row]
160
        Filter.__init__(self, filter_, reader)
161
162
class StripFilter(Filter):
163
    '''Strips whitespace'''
164
    def __init__(self, reader):
165
        def filter_(row): return [v.strip() for v in row]
166
        Filter.__init__(self, filter_, reader)
167 5570 aaronmk
168
class ColCtFilter(Filter):
169
    '''Gives all rows the same # columns'''
170
    def __init__(self, reader, cols_ct):
171
        def filter_(row): return util.list_as_length(row, cols_ct)
172
        Filter.__init__(self, filter_, reader)
173 5571 aaronmk
174
##### Translators
175
176 5586 aaronmk
class StreamFilter(Filter):
177
    '''Wraps a reader in a way that's usable to a filter stream that does not
178
    require lines to be strings. Reports EOF as '' instead of StopIteration.'''
179
    def __init__(self, reader):
180
        Filter.__init__(self, None, reader)
181
182
    def readline(self):
183
        try: return self.reader.next()
184
        except StopIteration: return '' # EOF
185
186 5735 aaronmk
class ColInsertFilter(Filter):
187 7290 aaronmk
    '''Adds column(s) to each row
188 9509 aaronmk
    @param mk_value(row, row_num) | literal_value
189 5735 aaronmk
    '''
190 7290 aaronmk
    def __init__(self, reader, mk_value, index=0, n=1):
191 9509 aaronmk
        if not callable(mk_value):
192
            value = mk_value
193
            def mk_value(row, row_num): return value
194
195 5735 aaronmk
        def filter_(row):
196
            row = list(row) # make sure it's mutable; don't modify input!
197 7290 aaronmk
            for i in xrange(n):
198
                row.insert(index+i, mk_value(row, self.reader.line_num))
199 5735 aaronmk
            return row
200
        Filter.__init__(self, filter_,
201
            streams.LineCountInputStream(StreamFilter(reader)))
202
203 5736 aaronmk
class RowNumFilter(ColInsertFilter):
204 5593 aaronmk
    '''Adds a row # column at the beginning of each row'''
205
    def __init__(self, reader):
206 5736 aaronmk
        def mk_value(row, row_num): return row_num
207
        ColInsertFilter.__init__(self, reader, mk_value, 0)
208 5593 aaronmk
209 5587 aaronmk
class InputRewriter(StreamFilter):
210 5571 aaronmk
    '''Wraps a reader, writing each row back to CSV'''
211
    def __init__(self, reader, dialect=csv.excel):
212 5587 aaronmk
        StreamFilter.__init__(self, reader)
213
214 5571 aaronmk
        self.dialect = dialect
215
216
    def readline(self):
217 8069 aaronmk
        try:
218
            row = self.reader.readline()
219
            if row == '': return row # EOF
220
221
            line_stream = StringIO.StringIO()
222
            csv.writer(line_stream, self.dialect).writerow(row)
223
            return line_stream.getvalue()
224
        except Exception, e:
225
            exc.print_ex(e)
226
            raise
227 5571 aaronmk
228
    def read(self, n): return self.readline() # forward all reads to readline()