From 99ff0a288c08ec2071139c6031eb880fa9c95310 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 14 Feb 2021 11:23:12 -0800 Subject: [PATCH] sqlite-utils insert --sniff option, closes #230 --- docs/cli.rst | 10 ++++++++-- sqlite_utils/cli.py | 28 +++++++++++++++++++++------- tests/sniff/example1.csv | 5 +++++ tests/sniff/example2.csv | 5 +++++ tests/sniff/example3.csv | 5 +++++ tests/sniff/example4.csv | 5 +++++ tests/test_sniff.py | 25 +++++++++++++++++++++++++ 7 files changed, 74 insertions(+), 9 deletions(-) create mode 100644 tests/sniff/example1.csv create mode 100644 tests/sniff/example2.csv create mode 100644 tests/sniff/example3.csv create mode 100644 tests/sniff/example4.csv create mode 100644 tests/test_sniff.py diff --git a/docs/cli.rst b/docs/cli.rst index 3ac5c4b..23b1690 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -498,7 +498,13 @@ A progress bar is displayed when inserting data from a file. You can hide the pr Alternative delimiters and quote characters ------------------------------------------- -If your file uses a delimiter other than ``,`` or a quote character other than ``"`` you can specify them using the ``--delimiter`` and ``--quotechar`` options. +If your file uses a delimiter other than ``,`` or a quote character other than ``"`` you can attempt to detect delimiters or you can specify them explicitly. + +The ``--sniff`` option can be used to attempt to detect the delimiters: + + sqlite-utils insert dogs.db dogs dogs.csv --sniff + +Alternatively, you can specify them using the ``--delimiter`` and ``--quotechar`` options. Here's a CSV file that uses ``;`` for delimiters and the ``|`` symbol for quote characters:: @@ -510,7 +516,7 @@ You can import that using:: $ sqlite-utils insert dogs.db dogs dogs.csv --delimiter=";" --quotechar="|" -Passing either ``--delimiter`` and ``--quotechar`` implies ``--csv``, so you can omit that option. +Passing ``--delimiter``, ``--quotechar`` or ``--sniff`` implies ``--csv``, so you can omit the ``--csv`` option. .. _cli_insert_replace: diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 3cb61af..c871232 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -603,6 +603,9 @@ def insert_upsert_options(fn): click.option("--tsv", is_flag=True, help="Expect TSV"), click.option("--delimiter", help="Delimiter to use for CSV files"), click.option("--quotechar", help="Quote character to use for CSV/TSV"), + click.option( + "--sniff", is_flag=True, help="Detect delimiter and quote character" + ), click.option( "--batch-size", type=int, default=100, help="Commit every X records" ), @@ -644,6 +647,7 @@ def insert_upsert_implementation( tsv, delimiter, quotechar, + sniff, batch_size, alter, upsert, @@ -658,33 +662,39 @@ def insert_upsert_implementation( ): db = sqlite_utils.Database(path) _load_extensions(db, load_extension) - if delimiter or quotechar: + if delimiter or quotechar or sniff: csv = True if (nl + csv + tsv) >= 2: raise click.ClickException("Use just one of --nl, --csv or --tsv") if encoding and not (csv or tsv): raise click.ClickException("--encoding must be used with --csv or --tsv") encoding = encoding or "utf-8" - json_file = io.TextIOWrapper(json_file, encoding=encoding) + buffered = io.BufferedReader(json_file, buffer_size=4096) + decoded = io.TextIOWrapper(buffered, encoding=encoding, line_buffering=True) if pk and len(pk) == 1: pk = pk[0] if csv or tsv: - dialect = "excel-tab" if tsv else "excel" - with file_progress(json_file, silent=silent) as json_file: + if sniff: + # Read first 2048 bytes and use that to detect + first_bytes = buffered.peek(2048) + dialect = csv_std.Sniffer().sniff(first_bytes.decode(encoding, "ignore")) + else: + dialect = "excel-tab" if tsv else "excel" + with file_progress(decoded, silent=silent) as decoded: csv_reader_args = {"dialect": dialect} if delimiter: csv_reader_args["delimiter"] = delimiter if quotechar: csv_reader_args["quotechar"] = quotechar - reader = csv_std.reader(json_file, **csv_reader_args) + reader = csv_std.reader(decoded, **csv_reader_args) headers = next(reader) docs = (dict(zip(headers, row)) for row in reader) else: try: if nl: - docs = (json.loads(line) for line in json_file) + docs = (json.loads(line) for line in decoded) else: - docs = json.load(json_file) + docs = json.load(decoded) if isinstance(docs, dict): docs = [docs] except json.decoder.JSONDecodeError: @@ -733,6 +743,7 @@ def insert( tsv, delimiter, quotechar, + sniff, batch_size, alter, encoding, @@ -761,6 +772,7 @@ def insert( tsv, delimiter, quotechar, + sniff, batch_size, alter=alter, upsert=False, @@ -790,6 +802,7 @@ def upsert( batch_size, delimiter, quotechar, + sniff, alter, not_null, default, @@ -813,6 +826,7 @@ def upsert( tsv, delimiter, quotechar, + sniff, batch_size, alter=alter, upsert=True, diff --git a/tests/sniff/example1.csv b/tests/sniff/example1.csv new file mode 100644 index 0000000..3daaadd --- /dev/null +++ b/tests/sniff/example1.csv @@ -0,0 +1,5 @@ +id,species,name,age +1,dog,Cleo,5 +2,dog,Pancakes,4 +3,cat,Mozie,8 +4,spider,"Daisy, the tarantula",6 diff --git a/tests/sniff/example2.csv b/tests/sniff/example2.csv new file mode 100644 index 0000000..0452e7f --- /dev/null +++ b/tests/sniff/example2.csv @@ -0,0 +1,5 @@ +id;species;name;age +1;dog;Cleo;5 +2;dog;Pancakes;4 +3;cat;Mozie;8 +4;spider;"Daisy, the tarantula";6 diff --git a/tests/sniff/example3.csv b/tests/sniff/example3.csv new file mode 100644 index 0000000..172c3d3 --- /dev/null +++ b/tests/sniff/example3.csv @@ -0,0 +1,5 @@ +id,species,name,age +1,dog,Cleo,5 +2,dog,Pancakes,4 +3,cat,Mozie,8 +4,spider,'Daisy, the tarantula',6 diff --git a/tests/sniff/example4.csv b/tests/sniff/example4.csv new file mode 100644 index 0000000..71b671e --- /dev/null +++ b/tests/sniff/example4.csv @@ -0,0 +1,5 @@ +id species name age +1 dog Cleo 5 +2 dog Pancakes 4 +3 cat Mozie 8 +4 spider 'Daisy, the tarantula' 6 diff --git a/tests/test_sniff.py b/tests/test_sniff.py new file mode 100644 index 0000000..36cc471 --- /dev/null +++ b/tests/test_sniff.py @@ -0,0 +1,25 @@ +from sqlite_utils import cli, Database +from click.testing import CliRunner +import pathlib +import pytest + +sniff_dir = pathlib.Path(__file__).parent / "sniff" + + +@pytest.mark.parametrize("filepath", sniff_dir.glob("example*")) +def test_sniff(tmpdir, filepath): + db_path = str(tmpdir / "test.db") + runner = CliRunner() + result = runner.invoke( + cli.cli, + ["insert", db_path, "creatures", str(filepath), "--sniff"], + catch_exceptions=False, + ) + assert result.exit_code == 0, result.stdout + db = Database(db_path) + assert list(db["creatures"].rows) == [ + {"id": "1", "species": "dog", "name": "Cleo", "age": "5"}, + {"id": "2", "species": "dog", "name": "Pancakes", "age": "4"}, + {"id": "3", "species": "cat", "name": "Mozie", "age": "8"}, + {"id": "4", "species": "spider", "name": "Daisy, the tarantula", "age": "6"}, + ]