From a03a4ead2cb9ea7fbac1c9592ea49ea16fce7474 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 24 Jul 2023 16:21:08 -0700
Subject: [PATCH] Support for --encoding, refs #8

---
 dclient/cli.py       |  8 +++++-
 tests/test_insert.py | 63 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/dclient/cli.py b/dclient/cli.py
index 4c7979e..49ff338 100644
--- a/dclient/cli.py
+++ b/dclient/cli.py
@@ -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
diff --git a/tests/test_insert.py b/tests/test_insert.py
index 1b521ab..096ad32 100644
--- a/tests/test_insert.py
+++ b/tests/test_insert.py
@@ -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