TSV and JSON support for sqlite-utils memory

Closes #281, closes #279, refs #272

sqlite_utils/cli.py

@@ -14,7 +14,14 @@ import os
 import sys
 import csv as csv_std
 import tabulate
-from .utils import file_progress, find_spatialite, sqlite3, decode_base64_values
+from .utils import (
+    file_progress,
+    find_spatialite,
+    sqlite3,
+    decode_base64_values,
+    rows_from_file,
+    Format,
+)
 
 CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
 
@@ -1175,18 +1182,21 @@ def memory(
         paths = [sql]
         sql = None
     for i, path in enumerate(paths):
-        if path == "-":
+        # Path may have a :format suffix
+        if ":" in path and path.rsplit(":", 1)[-1].upper() in Format.__members__:
+            path, suffix = path.rsplit(":", 1)
+            format = Format[suffix.upper()]
+        else:
+            format = None
+        if path in ("-", "stdin"):
             csv_fp = sys.stdin.buffer
             csv_table = "stdin"
         else:
             csv_path = pathlib.Path(path)
             csv_table = csv_path.stem
             csv_fp = csv_path.open("rb")
-
-        encoding = encoding or "utf-8-sig"
-        decoded_fp = io.TextIOWrapper(csv_fp, encoding=encoding)
-
-        db[csv_table].insert_all(csv_std.DictReader(decoded_fp))
+        rows = rows_from_file(csv_fp, format=format, encoding=encoding)
+        db[csv_table].insert_all(rows, alter=True)
         # Add convenient t / t1 / t2 views
         view_names = ["t{}".format(i + 1)]
         if i == 0:

sqlite_utils/utils.py

@@ -1,8 +1,13 @@
 import base64
-import click
 import contextlib
+import csv
+import enum
 import io
+import json
 import os
+from typing import Generator
+
+import click
 
 try:
     import pysqlite3 as sqlite3
@@ -111,3 +116,58 @@ def file_progress(file, silent=False, **kwargs):
         file_length = os.path.getsize(file.name)
         with click.progressbar(length=file_length, **kwargs) as bar:
             yield UpdateWrapper(file, bar.update)
+
+
+class Format(enum.Enum):
+    CSV = 1
+    TSV = 2
+    JSON = 3
+    NL = 4
+
+
+class RowsFromFileError(Exception):
+    pass
+
+
+class RowsFromFileBadJSON(RowsFromFileError):
+    pass
+
+
+def rows_from_file(
+    fp,
+    format=None,
+    dialect=None,
+    encoding=None,
+) -> Generator[dict, None, None]:
+    if format == Format.JSON:
+        decoded = json.load(fp)
+        if isinstance(decoded, dict):
+            decoded = [decoded]
+        if not isinstance(decoded, list):
+            raise RowsFromFileBadJSON("JSON must be a list or a dictionary")
+        yield from decoded
+    elif format == Format.NL:
+        yield from (json.loads(line) for line in fp if line.strip())
+    elif format == Format.CSV:
+        decoded_fp = io.TextIOWrapper(fp, encoding=encoding or "utf-8-sig")
+        yield from csv.DictReader(decoded_fp, dialect=dialect)
+    elif format == Format.TSV:
+        yield from rows_from_file(
+            fp, format=Format.CSV, dialect=csv.excel_tab, encoding=encoding
+        )
+    elif format is None:
+        # Detect the format, then call this recursively
+        buffered = io.BufferedReader(fp, buffer_size=4096)
+        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)
+        else:
+            dialect = csv.Sniffer().sniff(
+                first_bytes.decode(encoding or "utf-8-sig", "ignore")
+            )
+            yield from rows_from_file(
+                buffered, format=Format.CSV, dialect=dialect, encoding=encoding
+            )
+    else:
+        raise RowsFromFileError("Bad format")