diff --git a/docs/cli.rst b/docs/cli.rst index c6067ac..bad2d51 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -726,6 +726,32 @@ Data is expected to be encoded as Unicode UTF-8. If your data is an another char A progress bar is displayed when inserting data from a file. You can hide the progress bar using the ``--silent`` option. +By default every column inserted from a CSV or TSV file will be of type ``TEXT``. To automatically detect column types - resulting in a mix of ``TEXT``, ``INTEGER`` and ``FLOAT`` columns, use the ``--detect-types`` option (or its shortcut ``-d``). + +For example, given a ``creatures.csv`` file containing this:: + + name,age,weight + Cleo,6,45.5 + Dori,1,3.5 + +The following command:: + + $ sqlite-utils insert creatures.db creatures creatures.tsv --csv --detect-types + +Will produce this schema:: + + $ sqlite-utils schema creatures.db + CREATE TABLE "creatures" ( + [rowid] INTEGER PRIMARY KEY, + [name] TEXT, + [age] INTEGER, + [weight] FLOAT + ); + +You can set the ``SQLITE_UTILS_DETECT_TYPES`` environment variable if you want ``--detect-types`` to be the default behavior:: + + $ export SQLITE_UTILS_DETECT_TYPES=1 + .. _cli_insert_csv_tsv_delimiter: Alternative delimiters and quote characters diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index c5f50a1..96adaee 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -21,6 +21,7 @@ from .utils import ( decode_base64_values, rows_from_file, Format, + TypeTracker, ) CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"]) @@ -683,6 +684,13 @@ def insert_upsert_options(fn): "--encoding", help="Character encoding for input, defaults to utf-8", ), + click.option( + "-d", + "--detect-types", + is_flag=True, + envvar="SQLITE_UTILS_DETECT_TYPES", + help="Detect types for columns in CSV/TSV data", + ), load_extension_option, click.option("--silent", is_flag=True, help="Do not show progress bar"), ) @@ -712,6 +720,7 @@ def insert_upsert_implementation( not_null=None, default=None, encoding=None, + detect_types=None, load_extension=None, silent=False, ): @@ -728,6 +737,7 @@ def insert_upsert_implementation( encoding = encoding or "utf-8-sig" buffered = io.BufferedReader(json_file, buffer_size=4096) decoded = io.TextIOWrapper(buffered, encoding=encoding) + tracker = None if csv or tsv: if sniff: # Read first 2048 bytes and use that to detect @@ -749,6 +759,9 @@ def insert_upsert_implementation( else: headers = first_row docs = (dict(zip(headers, row)) for row in reader) + if detect_types: + tracker = TypeTracker() + docs = tracker.wrap(docs) else: try: if nl: @@ -781,6 +794,8 @@ def insert_upsert_implementation( "{}\n\nTry using --alter to add additional columns".format(e.args[0]) ) raise + if tracker is not None: + db[table].transform(types=tracker.types) @cli.command() @@ -815,6 +830,7 @@ def insert( batch_size, alter, encoding, + detect_types, load_extension, silent, ignore, @@ -849,6 +865,7 @@ def insert( replace=replace, truncate=truncate, encoding=encoding, + detect_types=detect_types, load_extension=load_extension, silent=silent, not_null=not_null, @@ -877,6 +894,7 @@ def upsert( not_null, default, encoding, + detect_types, load_extension, silent, ): diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 8a0b9ab..ffc727e 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -110,7 +110,16 @@ class UpdateWrapper: @contextlib.contextmanager def file_progress(file, silent=False, **kwargs): - if silent or file.fileno() == 0: # 0 = stdin + if silent: + yield file + return + # file.fileno() throws an exception in our test suite + try: + fileno = file.fileno() + except io.UnsupportedOperation: + yield file + return + if fileno == 0: # 0 means stdin yield file else: file_length = os.path.getsize(file.name) @@ -171,3 +180,68 @@ def rows_from_file( ) else: raise RowsFromFileError("Bad format") + + +class TypeTracker: + def __init__(self): + self.trackers = {} + + def wrap(self, iterator): + for row in iterator: + for key, value in row.items(): + tracker = self.trackers.setdefault(key, ValueTracker()) + tracker.evaluate(value) + yield row + + @property + def types(self): + return {key: tracker.guessed_type for key, tracker in self.trackers.items()} + + +class ValueTracker: + def __init__(self): + self.couldbe = {key: getattr(self, "test_" + key) for key in self.get_tests()} + + @classmethod + def get_tests(cls): + return [ + key.split("test_")[-1] + for key in cls.__dict__.keys() + if key.startswith("test_") + ] + + def test_integer(self, value): + try: + int(value) + return True + except ValueError: + return False + + def test_float(self, value): + try: + float(value) + return True + except ValueError: + return False + + def __repr__(self): + return self.guessed_type + ": possibilities = " + repr(self.couldbe) + + @property + def guessed_type(self): + options = set(self.couldbe.keys()) + # Return based on precedence + for key in self.get_tests(): + if key in options: + return key + return "text" + + def evaluate(self, value): + if not value or not self.couldbe: + return + not_these = [] + for name, test in self.couldbe.items(): + if not test(value): + not_these.append(name) + for key in not_these: + del self.couldbe[key] diff --git a/tests/test_cli.py b/tests/test_cli.py index b085ea4..8ddf71b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,6 +1,7 @@ from sqlite_utils import cli, Database from sqlite_utils.db import Index, ForeignKey from click.testing import CliRunner +from unittest import mock import json import os import pytest @@ -2067,3 +2068,33 @@ def test_csv_insert_bom(tmpdir): ("broken", "CREATE TABLE [broken] (\n [\ufeffname] TEXT,\n [age] TEXT\n)"), ("fixed", "CREATE TABLE [fixed] (\n [name] TEXT,\n [age] TEXT\n)"), ] + + +@pytest.mark.parametrize("option_or_env_var", (None, "-d", "--detect-types")) +def test_insert_detect_types(tmpdir, option_or_env_var): + db_path = str(tmpdir / "test.db") + data = "name,age,weight\nCleo,6,45.5\nDori,1,3.5" + extra = [] + if option_or_env_var: + extra = [option_or_env_var] + + def _test(): + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "creatures", "-", "--csv"] + extra, + catch_exceptions=False, + input=data, + ) + assert result.exit_code == 0 + db = Database(db_path) + assert list(db["creatures"].rows) == [ + {"rowid": 1, "name": "Cleo", "age": 6, "weight": 45.5}, + {"rowid": 2, "name": "Dori", "age": 1, "weight": 3.5}, + ] + + if option_or_env_var is None: + # Use environemnt variable instead of option + with mock.patch.dict(os.environ, {"SQLITE_UTILS_DETECT_TYPES": "1"}): + _test() + else: + _test()