Repository: ecmwf/earthkit-data
Pull request #277 (2024-01-15 12:32:57)

Implement multi part option for url streams 


Patch:
diff --git a/docs/guide/sources.rst b/docs/guide/sources.rst
index 17c7b20..e92bd63 100644
--- a/docs/guide/sources.rst
+++ b/docs/guide/sources.rst
@@ -152,7 +152,7 @@ url
   :param url: the URL to download
   :type url: str
   :param bool unpack: for archive formats such as ``.zip``, ``.tar``, ``.tar.gz``, etc, *earthkit-data* will attempt to open it and extract any usable file. To keep the downloaded file as is use ``unpack=False``
-  :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>`. This option only works for GRIB data. No archive formats supported (``unpack`` is ignored).
+  :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>`. This option only works for GRIB data. No archive formats supported (``unpack`` is ignored). ``stream`` only works for ``http`` and ``https`` URLs.
   :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 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
diff --git a/earthkit/data/sources/url.py b/earthkit/data/sources/url.py
index 4ae175d..4b328e7 100644
--- a/earthkit/data/sources/url.py
+++ b/earthkit/data/sources/url.py
@@ -22,6 +22,10 @@ from .file import FileSource
 LOG = logging.getLogger(__name__)
 
 
+def _ignore(*args, **kwargs):
+    pass
+
+
 def download_and_cache(
     url,
     *,
@@ -107,6 +111,7 @@ class UrlBase(FileSource):
     def __init__(
         self,
         url,
+        chunk_size=1024 * 1024,
         parts=None,
         filter=None,
         merger=None,
@@ -121,6 +126,7 @@ class UrlBase(FileSource):
         super().__init__(filter=filter, merger=merger)
 
         self.url = url
+        self.chunk_size = chunk_size
         self.parts = parts
         self.http_headers = http_headers
         self.auth = auth
@@ -181,7 +187,6 @@ class Url(UrlBase):
         self,
         url,
         *,
-        chunk_size=1024 * 1024,
         update_if_out_of_date=False,
         force=None,
         **kwargs,
@@ -208,7 +213,7 @@ class Url(UrlBase):
 
             self.downloader = Downloader(
                 self.url,
-                chunk_size=chunk_size,
+                chunk_size=self.chunk_size,
                 timeout=SETTINGS.get("url-download-timeout"),
                 verify=self.verify,
                 range_method=self.range_method,
@@ -295,6 +300,104 @@ class Url(UrlBase):
         return False
 
 
+class RequestIterStreamer:
+    """Expose fixed chunk-based stream reader used in mutiurl as a
+    stream supporting a generic read method
+    """
+
+    def __init__(self, iter_content):
+        from collections import deque
+
+        self.iter_content = iter_content
+        self.content = deque()
+        self.position = 0
+        self.total = 0
+        self.consumed = False
+
+    def _ensure_content(self, size):
+        while self.total < size:
+            try:
+                self.content.append(next(self.iter_content))
+                self.total += len(self.content[-1])
+            except StopIteration:
+                break
+
+    def _read(self, size):
+        assert len(self.content) > 0
+
+        start = self.position
+        length = min(len(self.content[0]) - start, size)
+        end = start + length
+        data = self.content[0][start:end]
+        last = 0
+        size -= length
+
+        if size > 0:
+            d = [data]
+            for i in range(1, len(self.content)):
+                start = 0
+                length = min(len(self.content[i]) - start, size)
+                end = start + length
+                d.append(self.content[i][start:end])
+                last = i
+                size -= length
+                if size <= 0:
+                    break
+            data = data = b"".join(d)
+
+        return data, last, end, size
+
+    def read(self, size):
+        if size <= 0 or self.consumed:
+            return bytes()
+
+        self._ensure_content(size)
+        if len(self.content) == 0 or self.total == 0:
+            self.close()
+            return bytes()
+
+        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:
+            if self.position == len(self.content[last]):
+                last += 1
+                self.position = 0
+
+            if last > 0:
+                for _ in range(0, last):
+                    self.content.popleft()
+
+            self.total = sum(len(x) for x in self.content)
+            self.total -= self.position
+
+        return data
+
+    def peek(self, size):
+        if size <= 0 or self.consumed:
+            return bytes()
+
+        self._ensure_content(size)
+        data, _, _, _ = self._read(size)
+        return data
+
+    def close(self):
+        if not self.closed:
+            self._clear()
+
+    @property
+    def closed(self):
+        return self.consumed
+
+    def _clear(self):
+        self.iter_content = None
+        self.content.clear()
+        self.position = 0
+        self.consumed = True
+
+
 class SingleUrlStream(UrlBase):
     def __init__(
         self,
@@ -306,50 +409,36 @@ class SingleUrlStream(UrlBase):
         if isinstance(self.url, (list, tuple)):
             raise TypeError("only a single url is supported")
 
+        from urllib.parse import urlparse
+
+        o = urlparse(self.url)
+        if o.scheme not in ("http", "https"):
+            raise NotImplementedError(f"streams are not supported for {o.scheme} URLs")
+
     def mutate(self):
         from .stream import _from_source
 
         return _from_source(self, **self._kwargs)
 
     def to_stream(self):
-        from urllib.request import Request, urlopen
-
-        headers = self.prepare_headers(self.url)
-
-        # TODO: ensure stream is closed when consumed
-        # TODO: use multiurl
-        r = Request(self.url, headers=headers)
-        return urlopen(r)
-
-    def prepare_headers(self, url):
-        headers = super().prepare_headers(url)
-        parts = self.parts_header(self.parts)
-        if parts is not None and parts:
-            if headers is None:
-                headers = {}
-            headers.update(parts)
-        return headers if headers is not None else {}
-
-    def parts_header(self, parts):
-        if parts is not None:
-            if isinstance(parts, (list, tuple)):
-                part = parts[0]
-            else:
-                part = parts
-
-            offset, length = part.offset, part.length
-            if offset is None:
-                offset = 0
-
-            start = offset
-            end = ""
-
-            if length is not None:
-                end = offset + length - 1
+        downloader = Downloader(
+            self.url,
+            chunk_size=256,
+            parts=self.parts,
+            timeout=SETTINGS.get("url-download-timeout"),
+            verify=self.verify,
+            range_method=self.range_method,
+            http_headers=self.prepare_headers(self.url),
+            fake_headers=self.fake_headers,
+            statistics_gatherer=_ignore,
+            progress_bar=progress_bar,
+            resume_transfers=False,
+            override_target_file=False,
+        )
 
-            return {"Range": f"bytes={start}-{end}"}
-        else:
-            return {}
+        size, mode, skip, trust_size = downloader.estimate_size(None)
+        stream = downloader.make_stream()
+        return RequestIterStreamer(stream(chunk_size=self.chunk_size))
 
 
 source = Url


Test patch:
diff --git a/tests/grib/test_grib_url_stream.py b/tests/grib/test_grib_url_stream.py
index 6367708..288fdef 100644
--- a/tests/grib/test_grib_url_stream.py
+++ b/tests/grib/test_grib_url_stream.py
@@ -361,17 +361,25 @@ def test_grib_multi_url_stream_memory():
 
 
 @pytest.mark.parametrize(
-    "parts,expected_meta",
+    "path,parts,expected_meta",
     [
-        ([(0, 150)], [("t", 1000)]),
-        ([(240, 150)], [("u", 1000)]),
-        ([(240, 480)], [("u", 1000), ("v", 1000)]),
+        ("examples/test6.grib", [(0, 150)], [("t", 1000)]),
+        ("examples/test6.grib", [(240, 150)], [("u", 1000)]),
+        ("examples/test6.grib", [(240, 480)], [("u", 1000), ("v", 1000)]),
+        ("test-data/karl_850.grib", [(0, 1683960)], [("t", 850)]),
+        ("test-data/karl_850.grib", [(0, 3367920)], [("t", 850), ("r", 850)]),
+        ("examples/test6.grib", [(240, 240), (720, 240)], [("u", 1000), ("t", 850)]),
+        (
+            "test-data/karl_850.grib",
+            [(0, 1683960), (3367920, 1683960)],
+            [("t", 850), ("z", 850)],
+        ),
     ],
 )
-def test_grib_single_url_stream_parts(parts, expected_meta):
+def test_grib_single_url_stream_parts(path, parts, expected_meta):
     ds = from_source(
         "url",
-        earthkit_remote_test_data_file("examples/test6.grib"),
+        earthkit_remote_test_data_file(path),
         parts=parts,
         stream=True,
     )
@@ -417,6 +425,11 @@ def test_grib_single_url_stream_parts(parts, expected_meta):
             [(0, 526)],
             [("u", 1000), ("2t", 0)],
         ),
+        (
+            [(240, 150), (720, 150)],
+            [(0, 526)],
+            [("u", 1000), ("t", 850), ("2t", 0)],
+        ),
     ],
 )
 def test_grib_multi_url_stream_parts(parts1, parts2, expected_meta):
@@ -444,6 +457,44 @@ def test_grib_multi_url_stream_parts(parts1, parts2, expected_meta):
     assert sum([1 for _ in ds]) == 0
 
 
+# Test RequestIterStreamer object from sources/url.py
+
+
+def iter_stream(chunk_size, data):
+    num = len(data)
+    pos = 0
+    while pos < num:
+        start = pos
+        end = min(pos + chunk_size, num)
+        pos += end - start
+        yield data[start:end]
+
+
+@pytest.mark.parametrize("chunk_size", [1, 2, 3, 4, 5, 10, 12, 14])
+@pytest.mark.parametrize("read_size", [1, 2, 3, 4, 5, 10, 12, 14])
+def test_request_iter_streamer(chunk_size, read_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(-1) == bytes()
+    assert stream.read(0) == bytes()
+    assert stream.peek(4) == data[:4]
+    assert not stream.closed
+
+    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(-1) == bytes()
+    assert stream.peek(4) == bytes()
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main