Project

General

Profile

1
# CSV I/O
2

    
3
import csv
4
import StringIO
5

    
6
import strings
7
import util
8

    
9
delims = ',\t`'
10
tsv_delim = '\t'
11
escape = '\\'
12

    
13
ending_placeholder = r'\n'
14

    
15
def is_tsv(dialect): return dialect.delimiter == tsv_delim
16

    
17
def sniff(line):
18
    '''Automatically detects the dialect'''
19
    dialect = csv.Sniffer().sniff(line, delims)
20
    # TSVs usually don't quote fields (nor doublequote embedded quotes)
21
    if is_tsv(dialect): dialect.quoting = csv.QUOTE_NONE
22
    else: dialect.doublequote = True # Sniffer doesn't turn this on by default
23
    return dialect
24

    
25
def stream_info(stream):
26
    '''Automatically detects the dialect based on the header line
27
    @return NamedTuple {header_line, dialect}'''
28
    info = util.NamedTuple()
29
    info.header_line = stream.readline()
30
    if info.header_line != '': info.dialect = sniff(info.header_line)
31
    else: info.dialect = None # line of '' indicates EOF = empty stream
32
    return info
33

    
34
class TsvReader:
35
    '''Unlike csv.reader, for TSVs, interprets \ as escaping a line ending but
36
    ignores it before everything else (e.g. \N for NULL). Also interprets the
37
    '\n' escape sequence as a newline.'''
38
    def __init__(self, stream, dialect):
39
        assert is_tsv(dialect)
40
        self.stream = stream
41
        self.dialect = dialect
42
    
43
    def __iter__(self): return self
44
    
45
    def next(self):
46
        record = ''
47
        ending = None
48
        while True:
49
            line = self.stream.readline()
50
            if line == '': raise StopIteration
51
            
52
            raw_contents, ending = strings.extract_line_ending(line)
53
            contents = strings.remove_suffix(escape, raw_contents)
54
            record += contents
55
            if len(contents) == len(raw_contents): break # no line continuation
56
            record += ending_placeholder
57
        row = csv.reader(StringIO.StringIO(record), self.dialect).next()
58
        return [v.replace(ending_placeholder, '\n') for v in row]
59

    
60
def reader_class(dialect):
61
    if is_tsv(dialect): return TsvReader
62
    else: return csv.reader
63

    
64
def make_reader(stream, dialect): return reader_class(dialect)(stream, dialect)
65

    
66
def reader_and_header(stream):
67
    '''Automatically detects the dialect based on the header line
68
    @return tuple (reader, header)'''
69
    info = stream_info(stream)
70
    reader_class_ = reader_class(info.dialect)
71
    header = reader_class_(StringIO.StringIO(info.header_line),
72
        info.dialect).next()
73
    return (reader_class_(stream, info.dialect), header)
74

    
75
##### csv modifications
76

    
77
# Note that these methods only work on *instances* of Dialect classes
78
csv.Dialect.__eq__ = lambda self, other: self.__dict__ == other.__dict__
79
csv.Dialect.__ne__ = lambda self, other: not (self == other)
(5-5/29)