TSV and JSON support for sqlite-utils memory

Closes #281, closes #279, refs #272

docs/cli.rst

@@ -234,21 +234,21 @@ This example attaches the ``books.db`` database under the alias ``books`` and th
     sqlite-utils dogs.db --attach books books.db \
        'select * from sqlite_master union all select * from books.sqlite_master'
 
-.. _cli_query_memory:
+.. _cli_memory:
 
-Querying CSV data directly using an in-memory database
+Querying data directly using an in-memory database
-======================================================
+==================================================
 
 The ``sqlite-utils memory`` command works similar to ``sqlite-utils query``, but allows you to execute queries against an in-memory database.
 
-You can also pass this command CSV files which will be loaded into a temporary in-memory table, allowing you to execute SQL against that data without a separate step to first convert it to SQLite.
+You can also pass this command CSV or JSON files which will be loaded into a temporary in-memory table, allowing you to execute SQL against that data without a separate step to first convert it to SQLite.
 
 Without any extra arguments, this command executes SQL against the in-memory database directly::
 
     $ sqlite-utils memory 'select sqlite_version()'
     [{"sqlite_version()": "3.35.5"}]
 
-It takes all of the same formatting options as :ref:`sqlite-utils query <cli_query>`: ``--csv`` and ``--csv`` and ``--table`` and ``--nl``::
+It takes all of the same output formatting options as :ref:`sqlite-utils query <cli_query>`: ``--csv`` and ``--csv`` and ``--table`` and ``--nl``::
 
     $ sqlite-utils memory 'select sqlite_version()' --csv             
     sqlite_version()
@@ -260,24 +260,28 @@ It takes all of the same formatting options as :ref:`sqlite-utils query <cli_que
     | 3.35.5             |
     +--------------------+
 
-.. _cli_query_memory_csv:
+.. _cli_memory_csv_json:
 
-Running queries directly against CSV
+Running queries directly against CSV or JSON
-------------------------------------
+--------------------------------------------
 
-If you have data in CSV format you can load it into an in-memory SQLite database and run queries against it directly in a single command using ``sqlite-utils memory`` like this::
+If you have data in CSV or JSON format you can load it into an in-memory SQLite database and run queries against it directly in a single command using ``sqlite-utils memory`` like this::
 
     $ sqlite-utils memory data.csv "select * from data"
 
-You can pass multiple files to the command if you want to run joins between different CSV files::
+You can pass multiple files to the command if you want to run joins between data from different files::
 
-    $ sqlite-utils memory one.csv two.csv "select * from one join two on one.id = two.other_id"
+    $ sqlite-utils memory one.csv two.json "select * from one join two on one.id = two.other_id"
 
-The in-memory tables will be named after the CSV files without their ``.csv`` extension. The tool also sets up aliases for those tables (using SQL views) as ``t1``, ``t2`` and so on, or you can use the alias ``t`` to refer to the first table::
+If your data is JSON it should be the same format supported by the :ref:`sqlite-utils insert command <cli_inserting_data>` - so either a single JSON object (treated as a single row) or a list of JSON objects.
+
+CSV data can be comma- or tab- delimited.
+
+The in-memory tables will be named after the files without their extensions. The tool also sets up aliases for those tables (using SQL views) as ``t1``, ``t2`` and so on, or you can use the alias ``t`` to refer to the first table::
 
     $ sqlite-utils memory example.csv "select * from t"
 
-To read from standard input, use ``-`` as the filename - then use ``stdin`` or ``t`` or ``t1`` as the table name::
+To read from standard input, use either ``-`` or ``stdin`` as the filename - then use ``stdin`` or ``t`` or ``t1`` as the table name::
 
     $ cat example.csv | sqlite-utils memory - "select * from stdin"
 
@@ -287,7 +291,24 @@ Incoming CSV data will be assumed to use ``utf-8``. If your data uses a differen
 
 If you are joining across multiple CSV files they must all use the same encoding.
 
-.. _cli_query_memory_attach:
+.. _cli_memory_explicit:
+
+Explicitly specifying the format
+--------------------------------
+
+By default, ``sqlite-utils memory`` will attempt to detect the incoming data format (JSON, TSV or CSV) automatically.
+
+You can instead specify an explicit format by adding a ``:csv``, ``:tsv``, ``:json`` or ``:nl`` (for newline-delimited JSON) suffix to the filename. For example::
+
+    $ sqlite-utils memory one.dat:csv two.dat:nl "select * from one union select * from two"
+
+Here the contents of ``one.dat`` will be treated as CSV and the contents of ``two.dat`` will be treated as newline-delimited JSON.
+
+To explicitly specify the format for data piped into the tool on standard input, use ``stdin:format`` - for example::
+
+    $ cat one.dat | sqlite-utils memory stdin:csv "select * from stdin"
+
+.. _cli_memory_attach:
 
 Joining in-memory data against existing databases using \-\-attach
 ------------------------------------------------------------------
@@ -303,7 +324,7 @@ Here the ``--attach trees trees.db`` option makes the ``trees.db`` database avai
 
 The CSV data that was piped into the script is available in the ``stdin`` table, so  ``... where rowid in (select id from stdin)`` can be used to return rows from the ``trees`` table that match IDs that were piped in as CSV content.
 
-.. _cli_query_memory_dump_save:
+.. _cli_memory_dump_save:
 
 \-\-dump and \-\-save
 ---------------------

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")
-
+        rows = rows_from_file(csv_fp, format=format, encoding=encoding)
-        encoding = encoding or "utf-8-sig"
+        db[csv_table].insert_all(rows, alter=True)
-        decoded_fp = io.TextIOWrapper(csv_fp, encoding=encoding)
-
-        db[csv_table].insert_all(csv_std.DictReader(decoded_fp))
         # 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")

tests/test_cli_memory.py

@@ -1,8 +1,10 @@
-from sqlite_utils import cli, Database
-from click.testing import CliRunner
-import pytest
 import json
 
+import pytest
+from click.testing import CliRunner
+
+from sqlite_utils import Database, cli
+
 
 def test_memory_basic():
     result = CliRunner().invoke(cli.cli, ["memory", "select 1 + 1"])
@@ -35,6 +37,78 @@ def test_memory_csv(tmpdir, sql_from, use_stdin):
     )
 
 
+@pytest.mark.parametrize("use_stdin", (True, False))
+def test_memory_tsv(tmpdir, use_stdin):
+    data = "id\tname\n1\tCleo\n2\tBants"
+    if use_stdin:
+        input = data
+        path = "stdin:tsv"
+        sql_from = "stdin"
+    else:
+        input = None
+        path = str(tmpdir / "chickens.tsv")
+        open(path, "w").write(data)
+        path = path + ":tsv"
+        sql_from = "chickens"
+    result = CliRunner().invoke(
+        cli.cli,
+        ["memory", path, "select * from {}".format(sql_from)],
+        input=data,
+    )
+    assert result.exit_code == 0, result.output
+    assert json.loads(result.output.strip()) == [
+        {"id": "1", "name": "Cleo"},
+        {"id": "2", "name": "Bants"},
+    ]
+
+
+@pytest.mark.parametrize("use_stdin", (True, False))
+def test_memory_json(tmpdir, use_stdin):
+    data = '[{"name": "Bants"}, {"name": "Dori", "age": 1}]'
+    if use_stdin:
+        input = data
+        path = "stdin:json"
+        sql_from = "stdin"
+    else:
+        input = None
+        path = str(tmpdir / "chickens.json")
+        open(path, "w").write(data)
+        path = path + ":json"
+        sql_from = "chickens"
+    result = CliRunner().invoke(
+        cli.cli,
+        ["memory", path, "select * from {}".format(sql_from)],
+        input=input,
+    )
+    assert result.exit_code == 0, result.output
+    assert json.loads(result.output.strip()) == [
+        {"name": "Bants", "age": None},
+        {"name": "Dori", "age": 1},
+    ]
+
+
+@pytest.mark.parametrize("use_stdin", (True, False))
+def test_memory_json_nl(tmpdir, use_stdin):
+    data = '{"name": "Bants"}\n\n{"name": "Dori"}'
+    if use_stdin:
+        input = data
+        path = "stdin:nl"
+        sql_from = "stdin"
+    else:
+        input = None
+        path = str(tmpdir / "chickens.json")
+        open(path, "w").write(data)
+        path = path + ":nl"
+        sql_from = "chickens"
+    result = CliRunner().invoke(
+        cli.cli,
+        ["memory", path, "select * from {}".format(sql_from)],
+        input=data,
+    )
+    assert result.exit_code == 0, result.output
+    assert json.loads(result.output.strip()) == [{"name": "Bants"}, {"name": "Dori"}]
+
+
 @pytest.mark.parametrize("use_stdin", (True, False))
 def test_memory_csv_encoding(tmpdir, use_stdin):
     latin1_csv = (