Repository: ecmwf/earthkit-data
Pull request #287 (2024-01-25 22:26:53)

Implement streams for polytope source


Patch:
diff --git a/docs/examples/polytope.ipynb b/docs/examples/polytope.ipynb
index 0108e52..ad7c100 100644
--- a/docs/examples/polytope.ipynb
+++ b/docs/examples/polytope.ipynb
@@ -8,6 +8,12 @@
     "## Retrieving data with polytope"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "3f96608d",
+   "metadata": {},
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "id": "75079b34-78ba-4536-8498-4dc3d0c3e646",
@@ -151,7 +157,7 @@
     "}\n",
     "\n",
     "\n",
-    "ds = earthkit.data.from_source(\"polytope\", \"ecmwf-mars\", request)\n",
+    "ds = earthkit.data.from_source(\"polytope\", \"ecmwf-mars\", request, stream=False)\n",
     "ds.ls()"
    ]
   },
diff --git a/docs/guide/sources.rst b/docs/guide/sources.rst
index a585ed9..09b7099 100644
--- a/docs/guide/sources.rst
+++ b/docs/guide/sources.rst
@@ -482,8 +482,8 @@ fdb
   The ``fdb`` source accesses the `FDB (Fields DataBase) <https://fields-database.readthedocs.io/en/latest/>`_, which is a domain-specific object store developed at ECMWF for storing, indexing and retrieving GRIB data. earthkit-data uses the `pyfdb <https://pyfdb.readthedocs.io/en/latest>`_ package to retrieve data from FDB.
 
   :param tuple *args: positional arguments specifying the request as a dict
-  :param bool stream: when it is ``True`` the data is read as a stream. Otherwise the data is retrieved into a file and stored in the :ref:`cache <caching>`.
-  :param int batch_size: used when ``stream=True`` and ``group_by`` is unset. It defines how many GRIB messages are consumed from the stream and kept in memory at a time. For details see :ref:`stream source <data-sources-stream>`.
+  :param bool stream: when it is ``True`` the data is read as a stream. Otherwise the data is retrieved into a file and stored in the :ref:`cache <caching>`. Only works for :ref:`grib` data.
+  :param int batch_size: used when ``stream=True`` and ``group_by`` is unset. It defines how many GRIB messages are consumed from the stream and kept in memory at a time. ``batch_size=0`` means all the data is read straight to memory. For details see :ref:`stream source <data-sources-stream>`.
   :param group_by: used when ``stream=True`` and can specify one or more metadata keys to control how GRIB messages are read from the stream. For details see :ref:`stream source <data-sources-stream>`.
   :type group_by: str, list of str
   :param dict **kwargs: other keyword arguments specifying the request
@@ -606,13 +606,17 @@ mars
 polytope
 --------
 
-.. py:function:: from_source("polytope", collection, *args, **kwargs)
+.. py:function:: from_source("polytope", collection, *args, stream=True,  batch_size=1, group_by=None, **kwargs)
   :noindex:
 
   The ``polytope`` source accesses the `Polytope web services <https://polytope-client.readthedocs.io/en/latest/>`_ , using the polytope-client_ package.
 
   :param str collection: the name of the polytope collection
   :param tuple *args: specify the request as a dict
+  :param bool stream: when it is ``True`` the data is read as a stream. Otherwise the data is retrieved into a file and stored in the :ref:`cache <caching>`. Only works for :ref:`grib` and CoverageJson data.
+  :param int batch_size: used when ``stream=True`` and ``group_by`` is unset. It defines how many GRIB messages are consumed from the stream and kept in memory at a time. ``batch_size=0`` means all the data is read straight to memory. For details see :ref:`stream source <data-sources-stream>`.
+  :param group_by: used when ``stream=True`` and can specify one or more metadata keys to control how GRIB messages are read from the stream. For details see :ref:`stream source <data-sources-stream>`.
+  :type group_by: str, list of str
   :param dict **kwargs: other keyword arguments specifying the request
 
   The following example retrieves GRIB data from the "ecmwf-mars" polytope collection:
@@ -635,7 +639,7 @@ polytope
           "domain": "g",
       }
 
-      ds = earthkit.data.from_source("polytope", "ecmwf-mars", request)
+      ds = earthkit.data.from_source("polytope", "ecmwf-mars", request, stream=False)
 
   Data downloaded from the polytope service is stored in the the :ref:`cache <caching>`. However,
   please note that, in the current version, each call to  :func:`from_source` will download the data again.
diff --git a/earthkit/data/readers/__init__.py b/earthkit/data/readers/__init__.py
index 2454554..8bd0944 100644
--- a/earthkit/data/readers/__init__.py
+++ b/earthkit/data/readers/__init__.py
@@ -188,7 +188,7 @@ def memory_reader(source, buffer):
     return _find_reader("memory_reader", source, buffer, magic=magic)
 
 
-def stream_reader(source, stream, memory):
+def stream_reader(source, stream, memory, content_type=None):
     """Create a reader for a stream"""
     magic = None
     if hasattr(stream, "peek") and callable(stream.peek):
@@ -206,4 +206,5 @@ def stream_reader(source, stream, memory):
         stream,
         magic=magic,
         memory=memory,
+        content_type=content_type,
     )
diff --git a/earthkit/data/readers/covjson.py b/earthkit/data/readers/covjson.py
index 5fbb71c..c908dae 100644
--- a/earthkit/data/readers/covjson.py
+++ b/earthkit/data/readers/covjson.py
@@ -13,7 +13,20 @@ import pathlib
 from earthkit.data.readers import Reader
 
 
-class CovjsonReader(Reader):
+class XarrayMixIn:
+    def to_xarray(self, **kwargs):
+        try:
+            from eccovjson.api import Eccovjson
+        except Exception:
+            raise ModuleNotFoundError(
+                "this feature requires 'eccovjson' to be installed!"
+            )
+
+        decoder = Eccovjson().decode(self._json())
+        return decoder.to_xarray()
+
+
+class CovjsonReader(XarrayMixIn, Reader):
     def __init__(self, source, path):
         super().__init__(source, path)
 
@@ -24,32 +37,118 @@ class CovjsonReader(Reader):
         # A Covjson is a source itself
         return self
 
-    def to_xarray(self, **kwargs):
-        try:
-            from eccovjson.api import Eccovjson
-        except Exception:
-            raise ModuleNotFoundError(
-                "this feature requires 'eccovjson' to be installed!"
-            )
-
+    def _json(self):
         import json
 
         with open(self.path, "r") as f:
             d = json.load(f)
-            decoder = Eccovjson().decode(d)
-            return decoder.to_xarray()
+            return d
+
+
+class CovjsonStreamReader(Reader):
+    def __init__(self, stream):
+        self.stream = stream
+
+    def __next__(self):
+        import json
+
+        d = self.stream.read()
+        if d:
+            return CovjsonInMemory(json.loads(d))
+        else:
+            raise StopIteration
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}"
+
+    def mutate_source(self):
+        # A Covjson is a source itself
+        return self
+
+
+class CovjsonMemoryReader(Reader):
+    def __init__(self, buf):
+        self.buf = buf
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}"
+
+    def mutate_source(self):
+        import json
+
+        return CovjsonInMemory(json.loads(self.buf))
+
+    @staticmethod
+    def _from_stream(stream):
+        d = stream.read()
+        return CovjsonMemoryReader(d)
+
+
+class CovjsonInMemory(XarrayMixIn, Reader):
+    def __init__(self, data):
+        self.data = data
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}"
+
+    def mutate_source(self):
+        # A Covjson is a source itself
+        return self
+
+    def _json(self):
+        return self.data
+
+
+def _match_content_type(content_type):
+    return content_type is not None and content_type == "application/prs.coverage+json"
+
+
+def _match_magic(magic, deeper_check):
+    if magic is not None:
+        type_id = b'{"type": "CoverageCollection"'
+        if not deeper_check:
+            return magic.startswith(type_id)
+        else:
+            return type_id in magic
+    return False
 
 
 def reader(
     source, path, *, magic=None, deeper_check=False, content_type=None, **kwargs
 ):
-    if content_type is not None and content_type == "application/prs.coverage+json":
+    def _reader():
         return CovjsonReader(source, path)
 
+    if _match_content_type(content_type) or _match_magic(magic, deeper_check):
+        return _reader()
+
     extension = pathlib.Path(path).suffix
     if extension in [".covjson"]:
-        return CovjsonReader(source, path)
+        return _reader()
 
     kind, compression = mimetypes.guess_type(path)
     if kind in ["application/prs.cov+json"]:
-        return CovjsonReader(source, path)
+        return _reader()
+
+
+def memory_reader(
+    source, buffer, *, magic=None, deeper_check=False, content_type=None, **kwargs
+):
+    if _match_content_type(content_type) or _match_magic(magic, deeper_check):
+        return CovjsonMemoryReader(buffer)
+
+
+def stream_reader(
+    source,
+    stream,
+    *,
+    magic=None,
+    deeper_check=False,
+    content_type=None,
+    memory=False,
+    **kwargs,
+):
+    if _match_content_type(content_type) or _match_magic(magic, deeper_check):
+        if memory:
+            return CovjsonMemoryReader._from_stream(stream)
+        return CovjsonStreamReader(stream)
diff --git a/earthkit/data/readers/grib/__init__.py b/earthkit/data/readers/grib/__init__.py
index 31aa6e5..c4e6586 100644
--- a/earthkit/data/readers/grib/__init__.py
+++ b/earthkit/data/readers/grib/__init__.py
@@ -12,7 +12,7 @@ LOG = logging.getLogger(__name__)
 
 
 def _match_magic(magic, deeper_check):
-    if magic is not None:
+    if magic is not None and len(magic) > 0:
         type_id = b"GRIB"
         if not deeper_check:
             return len(magic) >= 4 and magic[:4] == type_id
@@ -21,6 +21,12 @@ def _match_magic(magic, deeper_check):
     return False
 
 
+def _is_default(magic, content_type):
+    return (magic is None or len(magic) == 0) and (
+        content_type is None or len(content_type) == 0
+    )
+
+
 def reader(source, path, *, magic=None, deeper_check=False, **kwargs):
     if _match_magic(magic, deeper_check):
         from .reader import GRIBReader
@@ -36,10 +42,16 @@ def memory_reader(source, buffer, *, magic=None, deeper_check=False, **kwargs):
 
 
 def stream_reader(
-    source, stream, magic=None, *, deeper_check=False, memory=False, **kwargs
+    source,
+    stream,
+    magic=None,
+    *,
+    deeper_check=False,
+    content_type=None,
+    memory=False,
+    **kwargs
 ):
-    # by default we assume the stream is grib data
-    if magic is None or _match_magic(magic, deeper_check):
+    if _is_default(magic, content_type) or _match_magic(magic, deeper_check):
         from .memory import GribFieldListInMemory, GribStreamReader
 
         r = GribStreamReader(stream)
diff --git a/earthkit/data/sources/polytope.py b/earthkit/data/sources/polytope.py
index 9b6b29b..b4aa34d 100644
--- a/earthkit/data/sources/polytope.py
+++ b/earthkit/data/sources/polytope.py
@@ -12,6 +12,7 @@ import logging
 from . import Source
 from .multi_url import MultiUrl
 from .prompt import APIKeyPrompt
+from .url import Url
 
 LOG = logging.getLogger(__name__)
 
@@ -55,7 +56,7 @@ class Polytope(Source):
     >>> src.to_xarray()  # if datacube
     """
 
-    def __init__(self, dataset, request, address=None) -> None:
+    def __init__(self, dataset, request, address=None, stream=True, **kwargs) -> None:
         try:
             import polytope
         except ImportError:
@@ -66,6 +67,13 @@ class Polytope(Source):
         super().__init__()
         assert isinstance(dataset, str)
 
+        self._stream_kwargs = dict()
+        for k in ["group_by", "batch_size"]:
+            if k in kwargs:
+                self._stream_kwargs[k] = kwargs.pop(k)
+
+        self.stream = stream
+
         self.request = dict(dataset=dataset, request=request)
 
         credentials = PolytopeWebKeyPrompt().check(load=True)
@@ -88,7 +96,15 @@ class Polytope(Source):
 
         urls = [p["location"] for p in pointers]
         LOG.debug(f"{urls=}")
-        return MultiUrl(urls)
+
+        if self.stream:
+            return Url(
+                urls,
+                stream=True,
+                **self._stream_kwargs,
+            )
+        else:
+            return MultiUrl(urls)
 
 
 source = Polytope
diff --git a/earthkit/data/sources/stream.py b/earthkit/data/sources/stream.py
index d8f54f1..90e0146 100644
--- a/earthkit/data/sources/stream.py
+++ b/earthkit/data/sources/stream.py
@@ -269,7 +269,10 @@ def _from_source(source, **kwargs):
         raise TypeError(f"got invalid keyword argument(s): {list(kwargs.keys())}")
 
     if not isinstance(source, (list, tuple)):
-        maker = StreamSourceMaker(source, stream_kwargs)
+        source = [source]
+
+    if len(source) == 1:
+        maker = StreamSourceMaker(source[0], stream_kwargs)
         return maker()
     else:
         sources = [StreamSourceMaker(s, stream_kwargs) for s in source]
diff --git a/earthkit/data/sources/url.py b/earthkit/data/sources/url.py
index 016c8a3..c71dbd8 100644
--- a/earthkit/data/sources/url.py
+++ b/earthkit/data/sources/url.py
@@ -376,10 +376,13 @@ class RequestIterStreamer:
 
         return data, last, end, size
 
-    def read(self, size):
-        if size <= 0 or self.consumed:
+    def read(self, size=-1):
+        if size < -1 or size == 0 or self.consumed:
             return bytes()
 
+        if size == -1:
+            return self.readall()
+
         self._ensure_content(size)
         if len(self.content) == 0 or self.total == 0:
             self.close()
@@ -387,7 +390,6 @@ class RequestIterStreamer:
 
         data, last, self.position, missing_size = self._read(size)
         # LOG.debug(f"{size=} {last=} pos={self.position} {missing_size=}")
-
         if missing_size > 0:
             self.close()
         else:
@@ -404,6 +406,25 @@ class RequestIterStreamer:
 
         return data
 
+    def readall(self):
+        if self.consumed:
+            return bytes()
+
+        first = self.read(self.total)
+        if len(first) == 0:
+            first = next(self.iter_content)
+        res = [first]
+
+        for d in self.iter_content:
+            res.append(d)
+
+        self.close()
+
+        if len(res) == 1:
+            return res[0]
+        else:
+            return b"".join(res)
+
     def peek(self, size):
         if size <= 0 or self.consumed:
             return bytes()
@@ -452,7 +473,7 @@ class SingleUrlStream(UrlBase):
     def to_stream(self):
         downloader = Downloader(
             self.url,
-            chunk_size=256,
+            chunk_size=self.chunk_size,
             parts=self.parts,
             timeout=SETTINGS.get("url-download-timeout"),
             verify=self.verify,
@@ -466,6 +487,12 @@ class SingleUrlStream(UrlBase):
         )
 
         size, mode, skip, trust_size = downloader.estimate_size(None)
+
+        # cache data may contain the result of the http HEAD request
+        h = downloader.cache_data()
+        if isinstance(h, dict):
+            self.content_type = h.get("content-type")
+
         stream = downloader.make_stream()
         return RequestIterStreamer(stream(chunk_size=self.chunk_size))
 


Test patch:
diff --git a/tests/grib/test_grib_url_stream.py b/tests/grib/test_grib_url_stream.py
index 2c61ea1..ac5e109 100644
--- a/tests/grib/test_grib_url_stream.py
+++ b/tests/grib/test_grib_url_stream.py
@@ -522,7 +522,7 @@ def test_request_iter_streamer(chunk_size, read_size):
     stream = RequestIterStreamer(iter_stream(chunk_size, data))
 
     assert not stream.closed
-    assert stream.read(-1) == bytes()
+    assert stream.read(-2) == bytes()
     assert stream.read(0) == bytes()
     assert stream.peek(4) == data[:4]
     assert not stream.closed
@@ -530,6 +530,48 @@ def test_request_iter_streamer(chunk_size, read_size):
     for i in range(0, len(data), read_size):
         assert stream.read(read_size) == data[i : min(i + read_size, len(data))], i
 
+    assert stream.read(1) == bytes()
+    assert stream.closed
+    assert stream.read(0) == bytes()
+    assert stream.read() == bytes()
+    assert stream.read(-1) == bytes()
+    assert stream.peek(4) == bytes()
+
+
+@pytest.mark.parametrize("chunk_size", [1, 2, 3, 4, 5, 10, 12, 14])
+def test_request_iter_streamer_read_all_1(chunk_size):
+    from earthkit.data.sources.url import RequestIterStreamer
+
+    data = str.encode("0123456789abc")
+
+    stream = RequestIterStreamer(iter_stream(chunk_size, data))
+
+    assert not stream.closed
+    assert stream.read(-2) == bytes()
+    assert stream.read(0) == bytes()
+    assert stream.peek(4) == data[:4]
+    assert not stream.closed
+
+    assert stream.read() == data
+
+    assert stream.read(1) == bytes()
+    assert stream.closed
+    assert stream.read(0) == bytes()
+    assert stream.read(-1) == bytes()
+    assert stream.peek(4) == bytes()
+
+
+@pytest.mark.parametrize("chunk_size", [1, 2, 3, 4, 5, 10, 12, 14])
+def test_request_iter_streamer_read_all_2(chunk_size):
+    from earthkit.data.sources.url import RequestIterStreamer
+
+    data = str.encode("0123456789abc")
+
+    stream = RequestIterStreamer(iter_stream(chunk_size, data))
+
+    assert not stream.closed
+    assert stream.read() == data
+
     assert stream.read(1) == bytes()
     assert stream.closed
     assert stream.read(0) == bytes()
diff --git a/tests/readers/test_covjson_reader.py b/tests/readers/test_covjson_reader.py
index 1d6feb7..84a88c4 100644
--- a/tests/readers/test_covjson_reader.py
+++ b/tests/readers/test_covjson_reader.py
@@ -24,7 +24,42 @@ def test_covjson_to_xarray():
     ds = from_source("file", earthkit_test_data_file("time_series.covjson"))
     assert ds
     a = ds.to_xarray()
-    assert a
+    assert len(a.data_vars) == 1
+
+
+@pytest.mark.skipif(NO_ECCOVJSON, reason="no eccovjson available")
+def test_covjson_memory():
+    with open(earthkit_test_data_file("time_series.covjson"), "r") as f:
+        d = f.read().encode()
+
+    ds = from_source("memory", d)
+    assert ds
+    a = ds.to_xarray()
+    assert len(a.data_vars) == 1
+
+
+@pytest.mark.skipif(NO_ECCOVJSON, reason="no eccovjson available")
+def test_covjson_stream():
+    stream = open(earthkit_test_data_file("time_series.covjson"), "rb")
+
+    ds = from_source("stream", stream)
+    assert ds
+    c = next(ds)
+    a = c.to_xarray()
+    assert len(a.data_vars) == 1
+
+    with pytest.raises(StopIteration):
+        next(ds)
+
+
+@pytest.mark.skipif(NO_ECCOVJSON, reason="no eccovjson available")
+def test_covjson_stream_memory():
+    stream = open(earthkit_test_data_file("time_series.covjson"), "rb")
+
+    ds = from_source("stream", stream, batch_size=0)
+    assert ds
+    a = ds.to_xarray()
+    assert len(a.data_vars) == 1
 
 
 if __name__ == "__main__":
diff --git a/tests/sources/test_polytope.py b/tests/sources/test_polytope.py
index 039a953..7b9315f 100644
--- a/tests/sources/test_polytope.py
+++ b/tests/sources/test_polytope.py
@@ -45,7 +45,7 @@ def test_polytope_odb():
         "domain": "off",
     }
 
-    src = from_source("polytope", "ichange", request)
+    src = from_source("polytope", "ichange", request, stream=False)
     df = src.to_pandas()
     assert len(df) == 52
 
@@ -69,7 +69,7 @@ def test_polytope_grib():
         "domain": "g",
     }
 
-    ds = from_source("polytope", "ecmwf-mars", request)
+    ds = from_source("polytope", "ecmwf-mars", request, stream=False)
 
     assert len(ds) == 2
     assert ds.metadata("level") == [500, 500]