# CSV I/O

import csv

import StringIO

import strings

import util

delims = ',\t`'

tsv_delim = '\t'

escape = '\\'

ending_placeholder = r'\n'

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

def sniff(line):

    '''Automatically detects the dialect'''

    dialect = csv.Sniffer().sniff(line, delims)

    # TSVs usually don't quote fields (nor doublequote embedded quotes)

    if is_tsv(dialect): dialect.quoting = csv.QUOTE_NONE

    else: dialect.doublequote = True # Sniffer doesn't turn this on by default

    return dialect

def stream_info(stream, parse_header=False):

    '''Automatically detects the dialect based on the header line

    @return NamedTuple {header_line, header, dialect}'''

    info = util.NamedTuple()

    info.header_line = stream.readline()

    info.header = None

    if info.header_line != '':

        info.dialect = sniff(info.header_line)

        if parse_header:

            info.header = reader_class(info.dialect)(

                StringIO.StringIO(info.header_line), info.dialect).next()

    else: info.dialect = None # line of '' indicates EOF = empty stream

    return info

class TsvReader:

    '''Unlike csv.reader, for TSVs, interprets \ as escaping a line ending but

    ignores it before everything else (e.g. \N for NULL). Also interprets the

    '\n' escape sequence as a newline.'''

    def __init__(self, stream, dialect):

        assert is_tsv(dialect)

        self.stream = stream

        self.dialect = dialect

    def __iter__(self): return self

    def next(self):

        record = ''

        ending = None

        while True:

            line = self.stream.readline()

            if line == '': raise StopIteration

            raw_contents, ending = strings.extract_line_ending(line)

            contents = strings.remove_suffix(escape, raw_contents)

            record += contents

            if len(contents) == len(raw_contents): break # no line continuation

            record += ending_placeholder

        row = csv.reader(StringIO.StringIO(record), self.dialect).next()

        return [v.replace(ending_placeholder, '\n') for v in row]

def reader_class(dialect):

    if is_tsv(dialect): return TsvReader

    else: return csv.reader

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

def reader_and_header(stream):

    '''Automatically detects the dialect based on the header line

    @return tuple (reader, header)'''

    info = stream_info(stream, parse_header=True)

    return (make_reader(stream, info.dialect), info.header)

##### csv modifications

# Note that these methods only work on *instances* of Dialect classes

csv.Dialect.__eq__ = lambda self, other: self.__dict__ == other.__dict__

csv.Dialect.__ne__ = lambda self, other: not (self == other)

##### Row filters

class Filter:

    '''Wraps a reader, filtering each row'''

    def __init__(self, filter_, reader):

        self.reader = reader

        self.filter = filter_

    def __iter__(self): return self

    def next(self): return self.filter(self.reader.next())

std_nulls = [r'\N']

empty_nulls = [''] + std_nulls

class NullFilter(Filter):

    '''Translates special string values to None'''

    def __init__(self, reader, nulls=std_nulls):

        map_ = dict.fromkeys(nulls, None)

        def filter_(row): return [map_.get(v, v) for v in row]

        Filter.__init__(self, filter_, reader)

class StripFilter(Filter):

    '''Strips whitespace'''

    def __init__(self, reader):

        def filter_(row): return [v.strip() for v in row]

        Filter.__init__(self, filter_, reader)