Repository: intake/intake
Pull request #625 (2021-11-02 19:33:56)

Add drivers for .json and .jsonl file types


Patch:
diff --git a/docs/source/api_user.rst b/docs/source/api_user.rst
index 1b93ae1..874ab57 100644
--- a/docs/source/api_user.rst
+++ b/docs/source/api_user.rst
@@ -12,6 +12,8 @@ These are reference class and function definitions likely to be useful to everyo
    intake.open_
    intake.source.csv.CSVSource
    intake.source.textfiles.TextFilesSource
+   intake.source.jsonfiles.JSONFileSource
+   intake.source.jsonfiles.JSONLinesFileSource
    intake.source.npy.NPySource
    intake.source.zarr.ZarrArraySource
    intake.catalog.local.YAMLFileCatalog
@@ -51,6 +53,12 @@ Source classes
 .. autoclass:: intake.source.textfiles.TextFilesSource
    :members: __init__, discover, read_partition, read, to_dask, persist, export
 
+.. autoclass:: intake.source.jsonfiles.JSONFileSource
+   :members: __init__, discover, read, persist, export
+
+.. autoclass:: intake.source.jsonfiles.JSONLinesFileSource
+   :members: __init__, discover, read, head, persist, export
+
 .. autoclass:: intake.source.npy.NPySource
    :members: __init__, discover, read_partition, read, to_dask, persist, export
 
diff --git a/docs/source/plugin-directory.rst b/docs/source/plugin-directory.rst
index c51649b..0b22847 100644
--- a/docs/source/plugin-directory.rst
+++ b/docs/source/plugin-directory.rst
@@ -7,7 +7,8 @@ This is a list of known projects which install driver plugins for Intake, and th
 contains in parentheses:
 
 * builtin to Intake (``catalog``, ``csv``, ``intake_remote``, ``ndzarr``,
-  ``numpy``, ``textfiles``, ``yaml_file_cat``, ``yaml_files_cat``, ``zarr_cat``)
+  ``numpy``, ``textfiles``, ``yaml_file_cat``, ``yaml_files_cat``, ``zarr_cat``,
+  ``json``, ``jsonl``)
 * `intake-astro <https://github.com/intake/intake-astro>`_ Table and array loading of FITS astronomical data (``fits_array``, ``fits_table``)
 * `intake-accumulo <https://github.com/intake/intake-accumulo>`_ Apache Accumulo clustered data storage (``accumulo``)
 * `intake-avro <https://github.com/intake/intake-avro>`_: Apache Avro data serialization format (``avro_table``, ``avro_sequence``)
diff --git a/intake/source/jsonfiles.py b/intake/source/jsonfiles.py
new file mode 100644
index 0000000..ed5f665
--- /dev/null
+++ b/intake/source/jsonfiles.py
@@ -0,0 +1,168 @@
+import contextlib
+import json
+from itertools import islice
+
+from intake.source.base import DataSource
+
+
+class JSONFileSource(DataSource):
+    """
+    Read JSON files as a single dictionary or list
+
+    The files can be local or remote. Extra parameters for encoding, etc.,
+    go into ``storage_options``.
+    """
+
+    name = "json"
+    version = "0.0.1"
+    container = "python"
+
+    def __init__(
+        self,
+        urlpath: str,
+        text_mode: bool = True,
+        text_encoding: str = "utf8",
+        compression: str = None,
+        read: bool = True,
+        metadata: dict = None,
+        storage_options: dict = None,
+    ):
+        """
+        Parameters
+        ----------
+        urlpath : str
+            Target file. Can include protocol specified (e.g., "s3://").
+        text_mode : bool
+            Whether to open the file in text mode, recoding binary
+            characters on the fly
+        text_encoding : str
+            If text_mode is True, apply this encoding. UTF* is by far the most
+            common
+        compression : str or None
+            If given, decompress the file with the given codec on load. Can
+            be something like "zip", "gzip", "bz2", or to try to guess from the
+            filename, 'infer'
+        storage_options: dict
+            Options to pass to the file reader backend, including text-specific
+            encoding arguments, and parameters specific to the remote
+            file-system driver, if using.
+        """
+        from fsspec.utils import compressions
+
+        VALID_COMPRESSIONS = list(compressions.values()) + ["infer"]
+
+        self._urlpath = urlpath
+        self._storage_options = storage_options or {}
+        self._dataframe = None
+        self._file = None
+        self.compression = compression
+
+        if compression not in VALID_COMPRESSIONS:
+            raise ValueError(
+                f"Compression value {compression} must be one of {VALID_COMPRESSIONS}"
+            )
+        self.mode = "rt" if text_mode else "rb"
+        self.encoding = text_encoding
+        self._read = read
+
+        super(JSONFileSource, self).__init__(metadata=metadata)
+
+    def read(self):
+        import fsspec
+
+        urlpath = self._get_cache(self._urlpath)[0]
+        with fsspec.open(
+            urlpath,
+            mode=self.mode,
+            encoding=self.encoding,
+            compression=self.compression,
+            **self._storage_options,
+        ) as f:
+            return json.load(f)
+
+
+class JSONLinesFileSource(DataSource):
+    """
+    Read a JSONL (https://jsonlines.org/) file and return a list of objects,
+    each being valid json object (e.g. a dictionary or list)
+    """
+
+    name = "jsonl"
+    version = "0.0.1"
+    container = "python"
+
+    def __init__(
+        self,
+        urlpath: str,
+        text_mode: bool = True,
+        text_encoding: str = "utf8",
+        compression: str = None,
+        read: bool = True,
+        metadata: dict = None,
+        storage_options: dict = None,
+    ):
+        """
+        Parameters
+        ----------
+        urlpath : str
+            Target file. Can include protocol specified (e.g., "s3://").
+        text_mode : bool
+            Whether to open the file in text mode, recoding binary
+            characters on the fly
+        text_encoding : str
+            If text_mode is True, apply this encoding. UTF* is by far the most
+            common
+        compression : str or None
+            If given, decompress the file with the given codec on load. Can
+            be something like "zip", "gzip", "bz2", or to try to guess from the
+            filename, 'infer'.
+        storage_options: dict
+            Options to pass to the file reader backend, including text-specific
+            encoding arguments, and parameters specific to the remote
+            file-system driver, if using.
+        """
+        from fsspec.utils import compressions
+
+        VALID_COMPRESSIONS = list(compressions.values()) + ["infer"]
+
+        self._urlpath = urlpath
+        self._storage_options = storage_options or {}
+        self._dataframe = None
+        self._file = None
+        self.compression = compression
+        if compression not in VALID_COMPRESSIONS:
+            raise ValueError(
+                f"Compression value {compression} must be one of {VALID_COMPRESSIONS}"
+            )
+        self.mode = "rt" if text_mode else "rb"
+        self.encoding = text_encoding
+        self._read = read
+        super().__init__(metadata=metadata)
+
+    @contextlib.contextmanager
+    def _open(self):
+        """
+        Yields an fsspec.OpenFile object
+        """
+        import fsspec
+
+        urlpath = self._get_cache(self._urlpath)[0]
+        with fsspec.open(
+            urlpath,
+            mode=self.mode,
+            encoding=self.encoding,
+            compression=self.compression,
+            **self._storage_options,
+        ) as f:
+            yield f
+
+    def read(self):
+        with self._open() as f:
+            return list(map(json.loads, f))
+
+    def head(self, nrows: int = 100):
+        """
+        return the first `nrows` lines from the file
+        """
+        with self._open() as f:
+            return list(map(json.loads, islice(f, nrows)))
diff --git a/setup.py b/setup.py
index aeed12d..b909d4a 100644
--- a/setup.py
+++ b/setup.py
@@ -6,8 +6,10 @@
 # The full license is in the LICENSE file, distributed with this software.
 #-----------------------------------------------------------------------------
 
-from setuptools import setup, find_packages
 import sys
+
+from setuptools import find_packages, setup
+
 import versioneer
 
 requires = [line.strip() for line in open('requirements.txt').readlines()
@@ -49,6 +51,8 @@ setup(
             'yaml_files_cat = intake.catalog.local:YAMLFilesCatalog',
             'csv = intake.source.csv:CSVSource',
             'textfiles = intake.source.textfiles:TextFilesSource',
+            'json = intake.source.jsonfiles:JSONFileSource',
+            'jsonl = intake.source.jsonfiles:JSONLinesFileSource',
             'catalog = intake.catalog.base:Catalog',
             'intake_remote = intake.catalog.remote:RemoteCatalog',
             'numpy = intake.source.npy:NPySource',


Test patch:
diff --git a/intake/source/tests/test_json.py b/intake/source/tests/test_json.py
new file mode 100644
index 0000000..da641a9
--- /dev/null
+++ b/intake/source/tests/test_json.py
@@ -0,0 +1,70 @@
+# -----------------------------------------------------------------------------
+# Copyright (c) 2012 - 2021, Anaconda, Inc. and Intake contributors
+# All rights reserved.
+#
+# The full license is in the LICENSE file, distributed with this software.
+# -----------------------------------------------------------------------------
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, Optional
+
+import pytest
+from fsspec import open_files
+from fsspec.utils import compressions
+
+from intake.source.jsonfiles import JSONFileSource, JSONLinesFileSource
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+EXTENSIONS = {
+    compression: f".{extension}" for extension, compression in compressions.items()
+}
+
+
+@pytest.fixture(params=[None, "gzip", "bz2"])
+def json_file(request, tmp_path) -> str:
+    data = {"hello": "world"}
+    file_path = str(tmp_path / "1.json")
+    file_path += EXTENSIONS.get(request.param, "")
+    with open_files([file_path], mode="wt", compression=request.param)[0] as f:
+        f.write(json.dumps(data))
+    return file_path
+
+
+@pytest.fixture(params=[None, "gzip", "bz2"])
+def jsonl_file(request, tmp_path) -> str:
+    data = [{"hello": "world"}, [1, 2, 3]]
+    file_path = str(tmp_path / "1.jsonl")
+    file_path += EXTENSIONS.get(request.param, "")
+    with open_files([file_path], mode="wt", compression=request.param)[0] as f:
+        f.write("\n".join(json.dumps(row) for row in data))
+    return file_path
+
+
+def test_jsonfile(json_file: str):
+    j = JSONFileSource(json_file, text_mode=True, compression="infer")
+    out = j.read()
+    assert isinstance(out, dict)
+    assert out["hello"] == "world"
+
+
+def test_jsonlfile(jsonl_file: str):
+    j = JSONLinesFileSource(jsonl_file, compression="infer")
+    out = j.read()
+    assert isinstance(out, list)
+
+    assert isinstance(out[0], dict)
+    assert out[0]["hello"] == "world"
+
+    assert isinstance(out[1], list)
+    assert out[1] == [1, 2, 3]
+
+
+def test_jsonl_head(jsonl_file: str):
+    j = JSONLinesFileSource(jsonl_file, compression="infer")
+    out = j.head(1)
+    assert isinstance(out, list)
+    assert len(out) == 1
+    assert out[0]["hello"] == "world"