1# -*- coding: utf-8 -*-
2import warnings
3from .._utils import exhaustible
4from .._utils import seekable
5from .._utils import file_types
6from .get_reader import get_reader
7from .temptable import load_data
8from .temptable import savepoint
9
10
11preferred_encoding = 'utf-8'
12fallback_encoding = ['latin-1']
13
14
15def load_csv(cursor, table, csvfile, encoding=None, **kwds):
16    """Load *csvfile* and insert data into *table*."""
17    global preferred_encoding
18    global fallback_encoding
19
20    default = kwds.get('restval', '')  # Used for default column value.
21
22    if encoding:
23        # When an encoding is specified, use it to load *csvfile* or
24        # fail if there are errors (no fallback recovery):
25        with savepoint(cursor):
26            reader = get_reader.from_csv(csvfile, encoding, **kwds)
27            load_data(cursor, table, reader, default=default)
28
29        return  # <- EXIT!
30
31    # When the encoding is unspecified, try to load *csvfile* using the
32    # preferred encoding and failing that, try the fallback encodings:
33
34    if isinstance(csvfile, file_types) and seekable(csvfile):
35        position = csvfile.tell()  # Get current position if
36    else:                          # csvfile is file-like and
37        position = None            # supports random access.
38
39    try:
40        with savepoint(cursor):
41            reader = get_reader.from_csv(csvfile, preferred_encoding, **kwds)
42            load_data(cursor, table, reader, default=default)
43
44        return  # <- EXIT!
45
46    except UnicodeDecodeError as orig_error:
47        if exhaustible(csvfile) and position is None:
48            encoding, object_, start, end, reason = orig_error.args  # Unpack args.
49            reason = (
50                '{0}: unable to load {1!r}, cannot attempt fallback with '
51                '{2!r} type: must specify an appropriate text encoding'
52            ).format(reason, csvfile, csvfile.__class__.__name__)
53            raise UnicodeDecodeError(encoding, object_, start, end, reason)
54
55        if isinstance(fallback_encoding, list):
56            fallback_list = fallback_encoding
57        else:
58            fallback_list = [fallback_encoding]
59
60        for fallback in fallback_list:
61            if position is not None:
62                csvfile.seek(position)
63
64            try:
65                with savepoint(cursor):
66                    reader = get_reader.from_csv(csvfile, fallback, **kwds)
67                    load_data(cursor, table, reader, default=default)
68
69                msg = (
70                    '{0}: loaded {1!r} using fallback {2!r}: specify an '
71                    'appropriate text encoding to assure correct operation'
72                ).format(orig_error, csvfile, fallback)
73                warnings.warn(msg)
74
75                return  # <- EXIT!
76
77            except UnicodeDecodeError:
78                pass
79
80        # Note: DO NOT refactor this section using a for-else. I swear...
81        encoding, object_, start, end, reason = orig_error.args  # Unpack args.
82        reason = (
83            '{0}: unable to load {1!r}, fallback recovery unsuccessful: '
84            'must specify an appropriate text encoding'
85        ).format(reason, csvfile)
86        raise UnicodeDecodeError(encoding, object_, start, end, reason)
87