Project

General

Profile

1 1388 aaronmk
# CSV I/O
2
3
import csv
4
import StringIO
5
6 1623 aaronmk
import strings
7 1444 aaronmk
import util
8
9 1623 aaronmk
delims = ',\t`'
10
tsv_delim = '\t'
11
escape = '\\'
12 1442 aaronmk
13 1623 aaronmk
ending_placeholder = r'\n'
14
15
def is_tsv(dialect): return dialect.delimiter == tsv_delim
16
17 1621 aaronmk
def sniff(line):
18
    '''Automatically detects the dialect'''
19 1623 aaronmk
    dialect = csv.Sniffer().sniff(line, delims)
20 1621 aaronmk
    # TSVs usually don't quote fields (nor doublequote embedded quotes)
21 1623 aaronmk
    if is_tsv(dialect): dialect.quoting = csv.QUOTE_NONE
22 1621 aaronmk
    else: dialect.doublequote = True # Sniffer doesn't turn this on by default
23
    return dialect
24
25 1923 aaronmk
def stream_info(stream, parse_header=False):
26 1444 aaronmk
    '''Automatically detects the dialect based on the header line
27 1923 aaronmk
    @return NamedTuple {header_line, header, dialect}'''
28 1444 aaronmk
    info = util.NamedTuple()
29
    info.header_line = stream.readline()
30 1923 aaronmk
    info.header = None
31
    if info.header_line != '':
32
        info.dialect = sniff(info.header_line)
33
        if parse_header:
34
            info.header = reader_class(info.dialect)(
35
                StringIO.StringIO(info.header_line), info.dialect).next()
36 1660 aaronmk
    else: info.dialect = None # line of '' indicates EOF = empty stream
37 1444 aaronmk
    return info
38
39 1623 aaronmk
class TsvReader:
40
    '''Unlike csv.reader, for TSVs, interprets \ as escaping a line ending but
41
    ignores it before everything else (e.g. \N for NULL). Also interprets the
42
    '\n' escape sequence as a newline.'''
43
    def __init__(self, stream, dialect):
44
        assert is_tsv(dialect)
45
        self.stream = stream
46
        self.dialect = dialect
47
48
    def __iter__(self): return self
49
50
    def next(self):
51
        record = ''
52
        ending = None
53
        while True:
54
            line = self.stream.readline()
55
            if line == '': raise StopIteration
56
57
            raw_contents, ending = strings.extract_line_ending(line)
58
            contents = strings.remove_suffix(escape, raw_contents)
59
            record += contents
60
            if len(contents) == len(raw_contents): break # no line continuation
61
            record += ending_placeholder
62 3055 aaronmk
63
        # Prevent "new-line character seen in unquoted field" errors
64
        record = record.replace('\r', ending_placeholder)
65
66 1623 aaronmk
        row = csv.reader(StringIO.StringIO(record), self.dialect).next()
67
        return [v.replace(ending_placeholder, '\n') for v in row]
68
69
def reader_class(dialect):
70
    if is_tsv(dialect): return TsvReader
71
    else: return csv.reader
72
73
def make_reader(stream, dialect): return reader_class(dialect)(stream, dialect)
74
75 1388 aaronmk
def reader_and_header(stream):
76
    '''Automatically detects the dialect based on the header line
77
    @return tuple (reader, header)'''
78 1923 aaronmk
    info = stream_info(stream, parse_header=True)
79 1958 aaronmk
    return (make_reader(stream, info.dialect), info.header)
80 1446 aaronmk
81
##### csv modifications
82
83
# Note that these methods only work on *instances* of Dialect classes
84
csv.Dialect.__eq__ = lambda self, other: self.__dict__ == other.__dict__
85
csv.Dialect.__ne__ = lambda self, other: not (self == other)
86 2114 aaronmk
87
##### Row filters
88
89
class Filter:
90
    '''Wraps a reader, filtering each row'''
91
    def __init__(self, filter_, reader):
92
        self.reader = reader
93
        self.filter = filter_
94
95
    def __iter__(self): return self
96
97
    def next(self): return self.filter(self.reader.next())
98
99
std_nulls = [r'\N']
100
empty_nulls = [''] + std_nulls
101
102
class NullFilter(Filter):
103
    '''Translates special string values to None'''
104
    def __init__(self, reader, nulls=std_nulls):
105
        map_ = dict.fromkeys(nulls, None)
106
        def filter_(row): return [map_.get(v, v) for v in row]
107
        Filter.__init__(self, filter_, reader)
108
109
class StripFilter(Filter):
110
    '''Strips whitespace'''
111
    def __init__(self, reader):
112
        def filter_(row): return [v.strip() for v in row]
113
        Filter.__init__(self, filter_, reader)