Repository: ecmwf/earthkit-data
Pull request #276 (2024-01-11 16:32:38)

Implement parts for url stream
This PR makes parts (only single parts) work for url streams:

```python
# the remote file contains 6 GRIB messages. 
# We only read the first 2 (between bytes offset=0 and lenght=480).
ds = from_source(
        "url",
       "https://get.ecmwf.int/repository/test-data/earthkit-data/examples/test6.grib",
        parts=[(0, 480)],
        stream=True,
    )

for f in ds:
    # f is a field
    print(f)
```

output:
```
GribField(t,1000,20180801,1200,0,0)
GribField(u,1000,20180801,1200,0,0)
```


Patch:
diff --git a/earthkit/data/sources/multi_url.py b/earthkit/data/sources/multi_url.py
index f762303..f09caeb 100644
--- a/earthkit/data/sources/multi_url.py
+++ b/earthkit/data/sources/multi_url.py
@@ -13,7 +13,9 @@ from .multi import MultiSource
 
 
 class MultiUrl(MultiSource):
-    def __init__(self, urls, *args, filter=None, merger=None, force=None, **kwargs):
+    def __init__(
+        self, urls, *args, filter=None, merger=None, force=None, lazily=True, **kwargs
+    ):
         if not isinstance(urls, (list, tuple)):
             urls = [urls]
 
@@ -22,6 +24,9 @@ class MultiUrl(MultiSource):
 
         assert len(urls)
 
+        if len(urls) > 1 and isinstance(urls[0], str):
+            urls = sorted(urls)
+
         sources = [
             from_source(
                 "url",
@@ -30,9 +35,10 @@ class MultiUrl(MultiSource):
                 merger=merger,
                 force=force,
                 # Load lazily so we can do parallel downloads
-                lazily=True,
+                lazily=lazily,
+                **kwargs
             )
-            for url in sorted(urls)
+            for url in urls
         ]
 
         super().__init__(sources, filter=filter, merger=merger)
diff --git a/earthkit/data/sources/stream.py b/earthkit/data/sources/stream.py
index 4d460ed..d8f54f1 100644
--- a/earthkit/data/sources/stream.py
+++ b/earthkit/data/sources/stream.py
@@ -221,22 +221,21 @@ class StreamSource(StreamSourceBase):
     def mutate(self):
         assert self._reader_ is None
 
-        return make_source(
+        return _from_stream(
             self._stream, batch_size=self.batch_size, group_by=self.group_by
         )
 
 
 class StreamSourceMaker:
-    def __init__(self, create, data, stream_kwargs):
-        self.create = create
-        self.data = data
+    def __init__(self, source, stream_kwargs):
+        self.in_source = source
         self.stream_kwargs = dict(stream_kwargs)
         self.source = None
 
     def __call__(self):
         if self.source is None:
-            stream = self.create(self.data)
-            self.source = make_source(stream, **self.stream_kwargs)
+            stream = self.in_source.to_stream()
+            self.source = _from_stream(stream, **self.stream_kwargs)
 
             prev = None
             src = self.source
@@ -248,7 +247,7 @@ class StreamSourceMaker:
         return self.source
 
 
-def make_source(stream, group_by, batch_size):
+def _from_stream(stream, group_by, batch_size):
     _kwargs = dict(batch_size=batch_size, group_by=group_by)
 
     if group_by:
@@ -263,17 +262,17 @@ def make_source(stream, group_by, batch_size):
     raise ValueError(f"Unsupported stream parameters {batch_size=} {group_by=}")
 
 
-def make_stream_from_method(owner, create, data, **kwargs):
+def _from_source(source, **kwargs):
     stream_kwargs, kwargs = parse_stream_kwargs(**kwargs)
 
     if kwargs:
         raise TypeError(f"got invalid keyword argument(s): {list(kwargs.keys())}")
 
-    if not isinstance(data, (list, tuple)):
-        maker = StreamSourceMaker(create, data, stream_kwargs)
+    if not isinstance(source, (list, tuple)):
+        maker = StreamSourceMaker(source, stream_kwargs)
         return maker()
     else:
-        sources = [StreamSourceMaker(create, d, stream_kwargs) for d in data]
+        sources = [StreamSourceMaker(s, stream_kwargs) for s in source]
         return MultiStreamSource(sources, **stream_kwargs)
 
 
diff --git a/earthkit/data/sources/url.py b/earthkit/data/sources/url.py
index 6b60835..4ae175d 100644
--- a/earthkit/data/sources/url.py
+++ b/earthkit/data/sources/url.py
@@ -103,7 +103,7 @@ def download_and_cache(
     return path
 
 
-class Url(FileSource):
+class UrlBase(FileSource):
     def __init__(
         self,
         url,
@@ -111,42 +111,115 @@ class Url(FileSource):
         filter=None,
         merger=None,
         verify=True,
-        force=None,
-        chunk_size=1024 * 1024,
         range_method="auto",
         http_headers=None,
-        update_if_out_of_date=False,
         fake_headers=None,  # When HEAD is not allowed but you know the size
         stream=False,
+        auth=None,
         **kwargs,
     ):
         super().__init__(filter=filter, merger=merger)
 
-        # TODO: re-enable this feature
-        extension = None
-
         self.url = url
         self.parts = parts
+        self.http_headers = http_headers
+        self.auth = auth
+        self.verify = verify
+        self.range_method = range_method
+        self.fake_headers = fake_headers
         self.stream = stream
-        LOG.debug("URL %s", url)
+        self._kwargs = kwargs
+        LOG.debug(
+            f"url={self.url} parts={self.parts} auth={self.auth} _kwargs={self._kwargs}"
+        )
+
+    def connect_to_mirror(self, mirror):
+        return mirror.connection_for_url(self, self.url, self.parts)
+
+    def prepare_headers(self, url):
+        headers = {}
+        if self.http_headers is not None:
+            headers = dict(self.http_headers)
+
+        if self.auth is not None:
+            headers.update(self.auth.auth_header(url))
+
+        if not headers:
+            headers = None
+
+        return headers
+
+    def _get_parts(self, url):
+        if isinstance(self.parts, dict):
+            return self.parts.get(url, None)
+        else:
+            return self.parts
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.url})"
+
+
+class Url(UrlBase):
+    """Represent a URL source.
+
+    Parameters
+    ----------
+    url: str, list, tuple
+        Single url or a list/tuple of urls. A url item can be:
+
+        * a str
+        * a list/tuple of two items. The first item is the url as a str, while the second
+          item defines the parts for the given url. In this case the ``parts`` kwargs
+          cannot be used.
+
+        When ``url`` is a list/tuple these two formats cannot be mixed so either all the url
+        items are str or a list/tuple of url and parts.
+
+    """
+
+    def __init__(
+        self,
+        url,
+        *,
+        chunk_size=1024 * 1024,
+        update_if_out_of_date=False,
+        force=None,
+        **kwargs,
+    ):
+        super().__init__(url, **kwargs)
+
+        # TODO: re-enable this feature
+        extension = None
 
         if not self.stream:
             self.update_if_out_of_date = update_if_out_of_date
 
+            # ensure no parts kwargs is used when the parts are defined together with the urls
+            _parts_kwargs = {}
+            if not isinstance(url, (list, tuple)):
+                url = [url]
+            if isinstance(url[0], (list, tuple)):
+                if self.parts is not None:
+                    raise ValueError("cannot specify parts both as arg and kwarg")
+            else:
+                _parts_kwargs = {"parts": self.parts}
+
+            LOG.debug(f"http_headers={self.http_headers}")
+
             self.downloader = Downloader(
-                url,
+                self.url,
                 chunk_size=chunk_size,
                 timeout=SETTINGS.get("url-download-timeout"),
-                verify=verify,
-                parts=parts,
-                range_method=range_method,
-                http_headers=http_headers,
-                fake_headers=fake_headers,
+                verify=self.verify,
+                range_method=self.range_method,
+                http_headers=self.prepare_headers(self.url),
+                fake_headers=self.fake_headers,
                 statistics_gatherer=record_statistics,
                 progress_bar=progress_bar,
                 resume_transfers=True,
                 override_target_file=False,
                 download_file_extension=".download",
+                **_parts_kwargs,
             )
 
             if extension and extension[0] != ".":
@@ -166,37 +239,43 @@ class Url(FileSource):
                 self.downloader.download(target)
                 return self.downloader.cache_data()
 
-            print(f"{url=}")
-
             self.path = self.cache_file(
                 download,
-                dict(url=url, parts=parts),
+                dict(url=self.url, parts=self.parts),
                 extension=extension,
                 force=force,
             )
-        else:
-            self._kwargs = kwargs
 
     def mutate(self):
         if self.stream:
+            # create one stream source per url
+            from multiurl.downloader import _canonicalize
+
+            s = []
+            _kwargs = {}
+            if self.parts is not None:
+                _kwargs = {"parts": self.parts}
+            urls, _ = _canonicalize(self.url, **_kwargs)
+
+            for url, parts in urls:
+                s.append(
+                    SingleUrlStream(
+                        url,
+                        parts=parts,
+                        verify=True,
+                        range_method="auto",
+                        http_headers=self.prepare_headers(url),
+                        fake_headers=None,
+                        auth=self.auth,
+                    )
+                )
 
-            def url_to_stream(url):
-                from urllib.request import urlopen
-
-                # TODO: ensure stream is closed when consumed
-                return urlopen(url)
-
-            from .stream import make_stream_from_method
+            from .stream import _from_source
 
-            return make_stream_from_method(
-                self, url_to_stream, self.url, **self._kwargs
-            )
+            return _from_source(s, **self._kwargs)
         else:
             return super().mutate()
 
-    def connect_to_mirror(self, mirror):
-        return mirror.connection_for_url(self, self.url, self.parts)
-
     def out_of_date(self, url, path, cache_data):
         if SETTINGS.get("check-out-of-date-urls") is False:
             return False
@@ -215,8 +294,62 @@ class Url(FileSource):
                 )
         return False
 
-    def __repr__(self) -> str:
-        return f"Url({self.url})"
+
+class SingleUrlStream(UrlBase):
+    def __init__(
+        self,
+        url,
+        **kwargs,
+    ):
+        super().__init__(url, **kwargs)
+
+        if isinstance(self.url, (list, tuple)):
+            raise TypeError("only a single url is supported")
+
+    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
+
+            return {"Range": f"bytes={start}-{end}"}
+        else:
+            return {}
 
 
 source = Url
diff --git a/earthkit/data/utils/url.py b/earthkit/data/utils/url.py
new file mode 100644
index 0000000..1d43c60
--- /dev/null
+++ b/earthkit/data/utils/url.py
@@ -0,0 +1,22 @@
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+
+from abc import ABCMeta, abstractmethod
+
+
+class HttpAuthenticator(metaclass=ABCMeta):
+    @abstractmethod
+    def auth_header(self, url):
+        pass
+
+
+class AnonymousHttpAuthenticator(metaclass=ABCMeta):
+    def auth_header(self, url):
+        return {}


Test patch:
diff --git a/tests/grib/test_grib_url.py b/tests/grib/test_grib_url.py
new file mode 100644
index 0000000..476a337
--- /dev/null
+++ b/tests/grib/test_grib_url.py
@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import pytest
+
+from earthkit.data import from_source
+from earthkit.data.testing import earthkit_remote_test_data_file
+
+
+@pytest.mark.parametrize(
+    "parts,expected_meta",
+    [
+        ([(0, 150)], [("t", 1000)]),
+        ([(240, 150)], [("u", 1000)]),
+        ([(240, 480)], [("u", 1000), ("v", 1000)]),
+    ],
+)
+def test_grib_single_url_parts(parts, expected_meta):
+    ds = from_source(
+        "url", earthkit_remote_test_data_file("examples/test6.grib"), parts=parts
+    )
+
+    assert len(ds) == len(expected_meta)
+
+    cnt = 0
+    for i, f in enumerate(ds):
+        assert f.metadata(("param", "level")) == expected_meta[i], i
+        cnt += 1
+
+    assert cnt == len(expected_meta)
+
+
+@pytest.mark.parametrize(
+    "parts1,parts2,expected_meta",
+    [
+        (
+            [(240, 150)],
+            None,
+            [("u", 1000), ("2t", 0), ("msl", 0)],
+        ),
+        (
+            None,
+            [(0, 526)],
+            [
+                ("t", 1000),
+                ("u", 1000),
+                ("v", 1000),
+                ("t", 850),
+                ("u", 850),
+                ("v", 850),
+                ("2t", 0),
+            ],
+        ),
+        (
+            [(240, 150)],
+            [(0, 526)],
+            [("u", 1000), ("2t", 0)],
+        ),
+    ],
+)
+def test_grib_multi_url_parts(parts1, parts2, expected_meta):
+    ds = from_source(
+        "url",
+        [
+            [earthkit_remote_test_data_file("examples/test6.grib"), parts1],
+            [earthkit_remote_test_data_file("examples/test.grib"), parts2],
+        ],
+    )
+
+    assert len(ds) == len(expected_meta)
+
+    cnt = 0
+    for i, f in enumerate(ds):
+        assert f.metadata(("param", "level")) == expected_meta[i], i
+        cnt += 1
+
+    assert cnt == len(expected_meta)
+
+
+if __name__ == "__main__":
+    from earthkit.data.testing import main
+
+    main()
diff --git a/tests/grib/test_grib_url_stream.py b/tests/grib/test_grib_url_stream.py
index abdd03a..6367708 100644
--- a/tests/grib/test_grib_url_stream.py
+++ b/tests/grib/test_grib_url_stream.py
@@ -66,10 +66,14 @@ def test_grib_url_stream_group_by(_kwargs):
         [("t", 1000), ("u", 1000), ("v", 1000)],
         [("t", 850), ("u", 850), ("v", 850)],
     ]
+    cnt = 0
     for i, f in enumerate(fs):
         assert len(f) == 3
         assert f.metadata(("param", "level")) == ref[i]
         assert f.to_fieldlist("numpy") is not f
+        cnt += 1
+
+    assert cnt == len(ref)
 
     # stream consumed, no data is available
     assert sum([1 for _ in fs]) == 0
@@ -102,9 +106,12 @@ def test_grib_url_stream_single_batch(_kwargs):
         ("u", 850),
         ("v", 850),
     ]
-
+    cnt = 0
     for i, f in enumerate(ds):
         assert f.metadata(("param", "level")) == ref[i], i
+        cnt += 1
+
+    assert cnt == len(ref)
 
     # stream consumed, no data is available
     assert sum([1 for _ in ds]) == 0
@@ -129,9 +136,13 @@ def test_grib_url_stream_multi_batch(_kwargs, expected_meta):
     with pytest.raises(TypeError):
         len(ds)
 
+    cnt = 0
     for i, f in enumerate(ds):
         assert len(f) == len(expected_meta[i])
         f.metadata("param") == expected_meta[i]
+        cnt += 1
+
+    assert cnt == len(expected_meta)
 
     # stream consumed, no data is available
     assert sum([1 for _ in ds]) == 0
@@ -225,9 +236,12 @@ def test_grib_multi_url_stream_single_batch(_kwargs):
         ("t", 850),
         ("z", 850),
     ]
-
+    cnt = 0
     for i, f in enumerate(ds):
         assert f.metadata(("param", "level")) == ref[i], i
+        cnt += 1
+
+    assert cnt == len(ref)
 
     # stream consumed, no data is available
     assert sum([1 for _ in ds]) == 0
@@ -256,9 +270,13 @@ def test_grib_multi_url_stream_batch(_kwargs, expected_meta):
     with pytest.raises(TypeError):
         len(ds)
 
+    cnt = 0
     for i, f in enumerate(ds):
         assert len(f) == len(expected_meta[i])
         f.metadata("param") == expected_meta[i]
+        cnt += 1
+
+    assert cnt == len(expected_meta)
 
     # stream consumed, no data is available
     assert sum([1 for _ in ds]) == 0
@@ -342,6 +360,90 @@ def test_grib_multi_url_stream_memory():
     ]
 
 
+@pytest.mark.parametrize(
+    "parts,expected_meta",
+    [
+        ([(0, 150)], [("t", 1000)]),
+        ([(240, 150)], [("u", 1000)]),
+        ([(240, 480)], [("u", 1000), ("v", 1000)]),
+    ],
+)
+def test_grib_single_url_stream_parts(parts, expected_meta):
+    ds = from_source(
+        "url",
+        earthkit_remote_test_data_file("examples/test6.grib"),
+        parts=parts,
+        stream=True,
+    )
+
+    # no fieldlist methods are available
+    with pytest.raises(TypeError):
+        len(ds)
+
+    cnt = 0
+    for i, f in enumerate(ds):
+        assert f.metadata(("param", "level")) == expected_meta[i], i
+        cnt += 1
+
+    assert cnt == len(expected_meta)
+
+    # stream consumed, no data is available
+    assert sum([1 for _ in ds]) == 0
+
+
+@pytest.mark.parametrize(
+    "parts1,parts2,expected_meta",
+    [
+        (
+            [(240, 150)],
+            None,
+            [("u", 1000), ("2t", 0), ("msl", 0)],
+        ),
+        (
+            None,
+            [(0, 526)],
+            [
+                ("t", 1000),
+                ("u", 1000),
+                ("v", 1000),
+                ("t", 850),
+                ("u", 850),
+                ("v", 850),
+                ("2t", 0),
+            ],
+        ),
+        (
+            [(240, 150)],
+            [(0, 526)],
+            [("u", 1000), ("2t", 0)],
+        ),
+    ],
+)
+def test_grib_multi_url_stream_parts(parts1, parts2, expected_meta):
+    ds = from_source(
+        "url",
+        [
+            [earthkit_remote_test_data_file("examples/test6.grib"), parts1],
+            [earthkit_remote_test_data_file("examples/test.grib"), parts2],
+        ],
+        stream=True,
+    )
+
+    # no fieldlist methods are available
+    with pytest.raises(TypeError):
+        len(ds)
+
+    cnt = 0
+    for i, f in enumerate(ds):
+        assert f.metadata(("param", "level")) == expected_meta[i], i
+        cnt += 1
+
+    assert cnt == len(expected_meta)
+
+    # stream consumed, no data is available
+    assert sum([1 for _ in ds]) == 0
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main
 
diff --git a/tests/sources/test_url.py b/tests/sources/test_url.py
index 0dddf5e..1eabbc1 100644
--- a/tests/sources/test_url.py
+++ b/tests/sources/test_url.py
@@ -145,6 +145,20 @@ def test_url_netcdf_source_save():
     assert os.path.exists(tmp.path)
 
 
+def test_multi_url_parts_invalid():
+    parts1 = [(240, 150)]
+    parts2 = [(0, 526)]
+    with pytest.raises(ValueError):
+        from_source(
+            "url",
+            [
+                [earthkit_remote_test_data_file("examples/test6.grib"), parts1],
+                [earthkit_remote_test_data_file("examples/test.grib"), parts2],
+            ],
+            parts=[(0, 240)],
+        )
+
+
 if __name__ == "__main__":
     test_part_url()
     # from earthkit.data.testing import main