sqlite-utils insert --detect-types option, refs #282

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

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,
 ):

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]

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