# CSV I/O

import csv

import _csv

import StringIO

import dicts

import exc

import lists

import streams

import strings

import util

delims = ',;\t|`'

tab_padded_delims = ['\t|\t']

tsv_delim = '\t'

escape = '\\'

ending_placeholder = r'\n'

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

def sniff(line):

    '''Automatically detects the dialect'''

    line, ending = strings.extract_line_ending(line)

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

    except _csv.Error, e:

        if exc.e_msg(e) == 'Could not determine delimiter': dialect = csv.excel

        else: raise

    if is_tsv(dialect):

        dialect.quoting = csv.QUOTE_NONE

        # Check multi-char delims using \t

        delim = strings.find_any(line, tab_padded_delims)

        if delim:

            dialect.delimiter = delim

            line_suffix = delim.rstrip('\t')

            if line.endswith(line_suffix): ending = line_suffix+ending

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

    dialect.lineterminator = ending

    return dialect

def has_unbalanced_quotes(str_): return str_.count('"') % 2 == 1 # odd # of "

def has_multiline_column(str_): return has_unbalanced_quotes(str_)

def stream_info(stream, parse_header=False):

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

    Uses the Excel dialect if the CSV file is empty.

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

    info = util.NamedTuple()

    info.header_line = stream.readline()

    if has_multiline_column(info.header_line): # 1st line not full header

        # assume it's a header-only csv with multiline columns

        info.header_line += ''.join(stream.readlines()) # use entire file

    info.header = None

    if info.header_line != '':

        info.dialect = sniff(info.header_line)

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

    if parse_header:

        try: info.header = reader_class(info.dialect)(

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

        except StopIteration: info.header = []

    return info

tsv_encode_map = strings.json_encode_map[:]

tsv_encode_map.append(('\t', r'\t'))

tsv_decode_map = strings.flip_map(tsv_encode_map)

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 expands tsv_encode_map escapes.

    '''

    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

            line = strings.remove_suffix(self.dialect.lineterminator, line)

            contents = strings.remove_suffix(escape, line)

            record += contents

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

            record += ending_placeholder

        # Prevent "new-line character seen in unquoted field" errors

        record = record.replace('\r', ending_placeholder)

        # Split line

        if record == '': row = [] # csv.reader would interpret as EOF

        elif len(self.dialect.delimiter) > 1: # multi-char delims

            row = record.split(self.dialect.delimiter)

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

        return [strings.replace_all(tsv_decode_map, v) 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)

def header(stream):

    '''fetches just the header line of a stream'''

    reader, header = reader_and_header(stream)

    return 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)

__Dialect__validate_orig = csv.Dialect._validate

def __Dialect__validate(self):

        try: __Dialect__validate_orig(self)

        except _csv.Error, e:

            if str(e) == '"delimiter" must be an 1-character string': pass # OK

            else: raise

csv.Dialect._validate = __Dialect__validate

##### 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())

    def close(self): pass # support using as a stream

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)

class ColCtFilter(Filter):

    '''Gives all rows the same # columns'''

    def __init__(self, reader, cols_ct):

        def filter_(row): return util.list_as_length(row, cols_ct)

        Filter.__init__(self, filter_, reader)

##### Translators

class StreamFilter(Filter):

    '''Wraps a reader in a way that's usable to a filter stream that does not

    require lines to be strings. Reports EOF as '' instead of StopIteration.'''

    def __init__(self, reader):

        Filter.__init__(self, None, reader)

    def readline(self):

        try: return self.reader.next()

        except StopIteration: return '' # EOF

class ProgressInputFilter(streams.ProgressInputStream): # is also a reader

    # ProgressInputStream extends StreamIter, so this can be used as a reader

    '''wraps a reader, reporting the # rows read every n rows and after the last

    row is read

    @param log the output stream for progress messages

    '''

    def __init__(self, reader, log, msg='Read %d row(s)', n=100):

        streams.ProgressInputStream.__init__(self, StreamFilter(reader), log,

            msg, n)

class ColInsertFilter(Filter):

    '''Adds column(s) to each row

    @param mk_value(row, row_num) | literal_value

    '''

    def __init__(self, reader, mk_value, index=0, n=1, col_names=None):

        line_num_skip = 0

        if col_names != None:

            col_names = lists.mk_seq(col_names)

            n = len(col_names)

            line_num_skip = 1

        if not callable(mk_value):

            value = mk_value

            def mk_value(row, row_num): return value

        def filter_(row):

            row = list(row) # make sure it's mutable; don't modify input!

            if self.is_header and col_names != None:

                values = col_names

                self.is_header = False

            else: values = n*[mk_value(row, self.reader.line_num-line_num_skip)]

            for i in xrange(len(values)): row.insert(index+i, values[i])

            return row

        Filter.__init__(self, filter_,

            streams.LineCountInputStream(StreamFilter(reader)))

        self.is_header = True

class RowNumFilter(ColInsertFilter):

    '''Adds a row # column at the beginning of each row'''

    def __init__(self, reader, col_name=None):

        def mk_value(row, row_num): return row_num

        ColInsertFilter.__init__(self, reader, mk_value, col_names=col_name)

class InputRewriter(StreamFilter): # is also a stream

    '''Wraps a reader, writing each row back to CSV'''

    def __init__(self, reader, dialect=csv.excel):

        StreamFilter.__init__(self, reader)

        self.dialect = dialect

    def readline(self):

        try:

            row = StreamFilter.readline(self) # translate EOF

            if row == '': return row # EOF

            line_stream = StringIO.StringIO()

            csv.writer(line_stream, self.dialect).writerow(row)

            return line_stream.getvalue()

        except Exception, e:

            exc.print_ex(e)

            raise

    def read(self, n): return self.readline() # forward all reads to readline()

def row_dict_to_list(dict_, col_order=[]):

    '''translates a CSV dict-based row to a list-based one

    @param dict_ {'col': 'value', __}

    @return (header, row) = (['col', __], ['value', __])

    '''

    dict_ = dict_.copy() # don't modify input!

    pairs = []

    for col in col_order: pairs.append((col, dict_.pop(col))) # ordered cols 1st

    pairs += sorted(dict_.items()) # then remaining cols in alphabetical order

    return (dicts.pair_keys(pairs), dicts.pair_values(pairs))

class JsonReader(Filter):

    '''reads parsed JSON data as row tuples

    @param json_data [{'col': 'value', __}, __]

    '''

    def __init__(self, json_data, col_order=[]):

        self.header = None

        def filter_(row_dict):

            header, row = row_dict_to_list(row_dict, col_order)

            if self.header == None: # 1st JSON row: header

                self.header = header

                self.next_row = row

                row = header

            elif self.next_row != None: # 1st JSON row: data

                row = self.next_row

                self.next_row = None

            else: # remaining JSON rows

                assert header == self.header # all rows must have same cols

            return row

        Filter.__init__(self, filter_, iter(json_data))

        self.next_row = None