Support for --encoding, refs #8

This commit is contained in:
Simon Willison 2023-07-24 16:21:08 -07:00
commit a03a4ead2c
2 changed files with 68 additions and 3 deletions

View file

@ -1,5 +1,6 @@
import click
import httpx
import io
import itertools
import json
import pathlib
@ -92,6 +93,7 @@ def query(url_or_alias, sql, token):
@click.option("format_tsv", "--tsv", is_flag=True, help="Input is TSV")
@click.option("format_json", "--json", is_flag=True, help="Input is JSON")
@click.option("format_nl", "--nl", is_flag=True, help="Input is newline-delimited JSON")
@click.option("--encoding", help="Character encoding for CSV/TSV")
@click.option(
"--no-detect-types", is_flag=True, help="Don't detect column types for CSV/TSV"
)
@ -119,6 +121,7 @@ def insert(
format_tsv,
format_json,
format_nl,
encoding,
no_detect_types,
replace,
ignore,
@ -170,7 +173,10 @@ def insert(
fp = sys.stdin.buffer
file_size = None
rows, format = rows_from_file(fp, format=format)
try:
rows, format = rows_from_file(fp, format=format, encoding=encoding)
except Exception as ex:
raise click.ClickException(str(ex))
if format in (Format.JSON, Format.NL):
# Disable progress bar - it can't handle these formats

View file

@ -9,6 +9,11 @@ import pathlib
import pytest
@pytest.fixture
def assert_all_responses_were_requested() -> bool:
return False
@pytest.fixture
def non_mocked_hosts():
# This ensures httpx-mock will not affect Datasette's own
@ -63,6 +68,12 @@ SIMPLE_JSON = json.dumps(
]
)
SIMPLE_JSON_NL = '{"a": 1, "b": 2, "c": 3}\n'
LATIN1_CSV = (
b"date,name,latitude,longitude\n"
b"2020-01-01,Barra da Lagoa,-27.574,-48.422\n"
b"2020-03-04,S\xe3o Paulo,-23.561,-46.645\n"
b"2020-04-05,Salta,-24.793:-65.408"
)
InsertTest = namedtuple(
@ -122,6 +133,52 @@ def make_format_test(content, arg):
should_error=False,
expected_table_json=[{"rowid": 1, "a": "1", "b": "2", "c": "3"}],
),
# --encoding - without it this should error:
InsertTest(
input_data=LATIN1_CSV,
cmd_args=["--no-detect-types", "--create", "--csv"],
table_exists=False,
expected_output="Inserting rows\n",
should_error=True,
expected_table_json=None,
),
# --encoding - with it this should work:
InsertTest(
input_data=LATIN1_CSV,
cmd_args=[
"--no-detect-types",
"--create",
"--encoding",
"latin-1",
"--csv",
],
table_exists=False,
expected_output="Inserting rows\n",
should_error=False,
expected_table_json=[
{
"rowid": 1,
"date": "2020-01-01",
"name": "Barra da Lagoa",
"latitude": "-27.574",
"longitude": "-48.422",
},
{
"rowid": 2,
"date": "2020-03-04",
"name": "São Paulo",
"latitude": "-23.561",
"longitude": "-46.645",
},
{
"rowid": 3,
"date": "2020-04-05",
"name": "Salta",
"latitude": "-24.793:-65.408",
"longitude": None,
},
],
),
# Existing table, conflicting pk
InsertTest(
input_data=SIMPLE_CSV,
@ -220,7 +277,10 @@ def test_insert_against_datasette(
httpx_mock.add_callback(custom_response)
path = pathlib.Path(tmpdir) / "data.txt"
path.write_text(input_data)
if isinstance(input_data, str):
path.write_text(input_data)
else:
path.write_bytes(input_data)
runner = CliRunner()
result = runner.invoke(
cli,
@ -233,7 +293,6 @@ def test_insert_against_datasette(
token,
]
+ cmd_args,
catch_exceptions=False,
)
if not should_error:
assert result.exit_code == 0