--no-headers option for sqlite-utils insert --csv, closes #228

This commit is contained in:
Simon Willison 2021-02-14 14:25:03 -08:00
commit 50d2096f5e
3 changed files with 62 additions and 2 deletions

View file

@ -518,6 +518,17 @@ You can import that using::
Passing ``--delimiter``, ``--quotechar`` or ``--sniff`` implies ``--csv``, so you can omit the ``--csv`` option.
.. _cli_insert_csv_tsv_no_header:
CSV files without a header row
------------------------------
The first row of any CSV or TSV file is expected to contain the names of the columns in that file.
If your file does not include this row, you can use the ``--no-headers`` option to specify that the tool should not use that fist row as headers.
If you do this, the table will be created with column names called ``untitled_1`` and ``untitled_2`` and so on. You can then rename them using the ``sqlite-utils transform ... --rename`` command, see :ref:`cli_transform_table`.
.. _cli_insert_replace:
Insert-replacing data

View file

@ -618,6 +618,9 @@ def insert_upsert_options(fn):
click.option(
"--sniff", is_flag=True, help="Detect delimiter and quote character"
),
click.option(
"--no-headers", is_flag=True, help="CSV file has no header row"
),
click.option(
"--batch-size", type=int, default=100, help="Commit every X records"
),
@ -660,6 +663,7 @@ def insert_upsert_implementation(
delimiter,
quotechar,
sniff,
no_headers,
batch_size,
alter,
upsert,
@ -674,7 +678,7 @@ def insert_upsert_implementation(
):
db = sqlite_utils.Database(path)
_load_extensions(db, load_extension)
if delimiter or quotechar or sniff:
if delimiter or quotechar or sniff or no_headers:
csv = True
if (nl + csv + tsv) >= 2:
raise click.ClickException("Use just one of --nl, --csv or --tsv")
@ -699,7 +703,12 @@ def insert_upsert_implementation(
if quotechar:
csv_reader_args["quotechar"] = quotechar
reader = csv_std.reader(decoded, **csv_reader_args)
headers = next(reader)
first_row = next(reader)
if no_headers:
headers = ["untitled_{}".format(i + 1) for i in range(len(first_row))]
reader = itertools.chain([first_row], reader)
else:
headers = first_row
docs = (dict(zip(headers, row)) for row in reader)
else:
try:
@ -756,6 +765,7 @@ def insert(
delimiter,
quotechar,
sniff,
no_headers,
batch_size,
alter,
encoding,
@ -785,6 +795,7 @@ def insert(
delimiter,
quotechar,
sniff,
no_headers,
batch_size,
alter=alter,
upsert=False,
@ -815,6 +826,7 @@ def upsert(
delimiter,
quotechar,
sniff,
no_headers,
alter,
not_null,
default,
@ -839,6 +851,7 @@ def upsert(
delimiter,
quotechar,
sniff,
no_headers,
batch_size,
alter=alter,
upsert=True,

View file

@ -1819,3 +1819,39 @@ def test_long_csv_column_value(tmpdir):
rows = list(db["bigtable"].rows)
assert len(rows) == 1
assert rows[0]["text"] == long_string
@pytest.mark.parametrize(
"args",
(
["--csv", "--no-headers"],
["--no-headers"],
),
)
def test_csv_import_no_headers(tmpdir, args):
db_path = str(tmpdir / "test.db")
csv_path = str(tmpdir / "test.csv")
csv_file = open(csv_path, "w")
csv_file.write("Cleo,Dog,5\n")
csv_file.write("Tracy,Spider,7\n")
csv_file.close()
result = CliRunner().invoke(
cli.cli,
["insert", db_path, "creatures", csv_path] + args,
catch_exceptions=False,
)
assert result.exit_code == 0
db = Database(db_path)
schema = db["creatures"].schema
assert schema == (
"CREATE TABLE [creatures] (\n"
" [untitled_1] TEXT,\n"
" [untitled_2] TEXT,\n"
" [untitled_3] TEXT\n"
")"
)
rows = list(db["creatures"].rows)
assert rows == [
{"untitled_1": "Cleo", "untitled_2": "Dog", "untitled_3": "5"},
{"untitled_1": "Tracy", "untitled_2": "Spider", "untitled_3": "7"},
]