Don't detect types on JSON input to memory, closes #283

This commit is contained in:
Simon Willison 2021-06-19 07:52:44 -07:00
commit dc94f4bb8c
3 changed files with 19 additions and 16 deletions

View file

@ -1241,9 +1241,9 @@ def memory(
csv_path = pathlib.Path(path)
csv_table = csv_path.stem
csv_fp = csv_path.open("rb")
rows = rows_from_file(csv_fp, format=format, encoding=encoding)
rows, format_used = rows_from_file(csv_fp, format=format, encoding=encoding)
tracker = None
if not no_detect_types:
if format_used in (Format.CSV, Format.TSV) and not no_detect_types:
tracker = TypeTracker()
rows = tracker.wrap(rows)
db[csv_table].insert_all(rows, alter=True)

View file

@ -155,15 +155,18 @@ def rows_from_file(
decoded = [decoded]
if not isinstance(decoded, list):
raise RowsFromFileBadJSON("JSON must be a list or a dictionary")
yield from decoded
return decoded, Format.JSON
elif format == Format.NL:
yield from (json.loads(line) for line in fp if line.strip())
return (json.loads(line) for line in fp if line.strip()), Format.NL
elif format == Format.CSV:
decoded_fp = io.TextIOWrapper(fp, encoding=encoding or "utf-8-sig")
yield from csv.DictReader(decoded_fp, dialect=dialect)
return csv.DictReader(decoded_fp, dialect=dialect), Format.CSV
elif format == Format.TSV:
yield from rows_from_file(
fp, format=Format.CSV, dialect=csv.excel_tab, encoding=encoding
return (
rows_from_file(
fp, format=Format.CSV, dialect=csv.excel_tab, encoding=encoding
)[0],
Format.TSV,
)
elif format is None:
# Detect the format, then call this recursively
@ -171,12 +174,12 @@ def rows_from_file(
first_bytes = buffered.peek(2048).strip()
if first_bytes.startswith(b"[") or first_bytes.startswith(b"{"):
# TODO: Detect newline-JSON
yield from rows_from_file(buffered, format=Format.JSON)
return rows_from_file(buffered, format=Format.JSON)
else:
dialect = csv.Sniffer().sniff(
first_bytes.decode(encoding or "utf-8-sig", "ignore")
)
yield from rows_from_file(
return rows_from_file(
buffered, format=Format.CSV, dialect=dialect, encoding=encoding
)
else:
@ -215,14 +218,14 @@ class ValueTracker:
try:
int(value)
return True
except ValueError:
except (ValueError, TypeError):
return False
def test_float(self, value):
try:
float(value)
return True
except ValueError:
except (ValueError, TypeError):
return False
def __repr__(self):

View file

@ -64,7 +64,7 @@ def test_memory_tsv(tmpdir, use_stdin):
@pytest.mark.parametrize("use_stdin", (True, False))
def test_memory_json(tmpdir, use_stdin):
data = '[{"name": "Bants"}, {"name": "Dori", "age": 1}]'
data = '[{"name": "Bants"}, {"name": "Dori", "age": 1, "nested": {"nest": 1}}]'
if use_stdin:
input = data
path = "stdin:json"
@ -82,8 +82,8 @@ def test_memory_json(tmpdir, use_stdin):
)
assert result.exit_code == 0, result.output
assert json.loads(result.output.strip()) == [
{"rowid": 1, "name": "Bants", "age": None},
{"rowid": 2, "name": "Dori", "age": 1},
{"name": "Bants", "age": None, "nested": None},
{"name": "Dori", "age": 1, "nested": '{"nest": 1}'},
]
@ -107,8 +107,8 @@ def test_memory_json_nl(tmpdir, use_stdin):
)
assert result.exit_code == 0, result.output
assert json.loads(result.output.strip()) == [
{"rowid": 1, "name": "Bants"},
{"rowid": 2, "name": "Dori"},
{"name": "Bants"},
{"name": "Dori"},
]