implement HTTPClient.download and add tests

Author: geetu040

openml/_api/clients/http.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
+import hashlib
 import json
 import logging
 import math
 import random
 import time
 import xml
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from pathlib import Path
 from typing import Any
 from urllib.parse import urlencode, urljoin, urlparse
@@ -18,6 +19,8 @@
 from openml.__version__ import __version__
 from openml.enums import RetryPolicy
 from openml.exceptions import (
+    OpenMLCacheRequiredError,
+    OpenMLHashException,
     OpenMLNotAuthorizedError,
     OpenMLServerError,
     OpenMLServerException,
@@ -315,14 +318,15 @@ def _request(  # noqa: PLR0913
 
         return response, retry_raise_e
 
-    def request(
+    def request(  # noqa: PLR0913, C901
         self,
         method: str,
         path: str,
         *,
         use_cache: bool = False,
         reset_cache: bool = False,
         use_api_key: bool = False,
+        md5_checksum: str | None = None,
         **request_kwargs: Any,
     ) -> Response:
         url = urljoin(self.server, urljoin(self.base_url, path))
@@ -384,15 +388,28 @@ def request(
             cache_key = self.cache.get_key(url, params)
             self.cache.save(cache_key, response)
 
+        if md5_checksum is not None:
+            self._verify_checksum(response, md5_checksum)
+
         return response
 
+    def _verify_checksum(self, response: Response, md5_checksum: str) -> None:
+        # ruff sees hashlib.md5 as insecure
+        actual = hashlib.md5(response.content).hexdigest()  # noqa: S324
+        if actual != md5_checksum:
+            raise OpenMLHashException(
+                f"Checksum of downloaded file is unequal to the expected checksum {md5_checksum} "
+                f"when downloading {response.url}.",
+            )
+
     def get(
         self,
         path: str,
         *,
         use_cache: bool = False,
         reset_cache: bool = False,
         use_api_key: bool = False,
+        md5_checksum: str | None = None,
         **request_kwargs: Any,
     ) -> Response:
         return self.request(
@@ -401,19 +418,22 @@ def get(
             use_cache=use_cache,
             reset_cache=reset_cache,
             use_api_key=use_api_key,
+            md5_checksum=md5_checksum,
             **request_kwargs,
         )
 
     def post(
         self,
         path: str,
+        *,
+        use_api_key: bool = True,
         **request_kwargs: Any,
     ) -> Response:
         return self.request(
             method="POST",
             path=path,
             use_cache=False,
-            use_api_key=True,
+            use_api_key=use_api_key,
             **request_kwargs,
         )
 
@@ -429,3 +449,33 @@ def delete(
             use_api_key=True,
             **request_kwargs,
         )
+
+    def download(
+        self,
+        url: str,
+        handler: Callable[[Response, Path, str], Path] | None = None,
+        encoding: str = "utf-8",
+        file_name: str = "response.txt",
+        md5_checksum: str | None = None,
+    ) -> Path:
+        if self.cache is None:
+            raise OpenMLCacheRequiredError(
+                "A cache object is required for download, but none was provided in the HTTPClient."
+            )
+        base = self.cache.path
+        file_path = base / "downloads" / urlparse(url).path.lstrip("/") / file_name
+        file_path = file_path.expanduser()
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        if file_path.exists():
+            return file_path
+
+        response = self.get(url, md5_checksum=md5_checksum)
+        if handler is not None:
+            return handler(response, file_path, encoding)
+
+        return self._text_handler(response, file_path, encoding)
+
+    def _text_handler(self, response: Response, path: Path, encoding: str) -> Path:
+        with path.open("w", encoding=encoding) as f:
+            f.write(response.text)
+        return path

openml/exceptions.py

@@ -69,3 +69,7 @@ class ObjectNotPublishedError(PyOpenMLError):
 
 class OpenMLNotSupportedError(PyOpenMLError):
     """Raised when an API operation is not supported for a resource/version."""
+
+
+class OpenMLCacheRequiredError(PyOpenMLError):
+    """Raised when a cache object is required but not provided."""

tests/test_api/test_http.py

@@ -4,9 +4,11 @@
 import pytest
 from openml.testing import TestAPIBase
 import os
+from pathlib import Path
 from urllib.parse import urljoin
 from openml.enums import APIVersion
 from openml._api import HTTPClient
+from openml.exceptions import OpenMLCacheRequiredError
 
 
 class TestHTTPClient(TestAPIBase):
@@ -174,3 +176,67 @@ def test_post_and_delete(self):
             if task_id is not None:
                 del_response = self.http_client.delete(f"task/{task_id}")
                 self.assertEqual(del_response.status_code, 200)
+
+    def test_download_requires_cache(self):
+        client = HTTPClient(
+            server=self.http_client.server,
+            base_url=self.http_client.base_url,
+            api_key=self.http_client.api_key,
+            retries=1,
+            retry_policy=self.http_client.retry_policy,
+            cache=None,
+        )
+
+        with pytest.raises(OpenMLCacheRequiredError):
+            client.download("https://www.openml.org")
+
+    @pytest.mark.uses_test_server()
+    def test_download_creates_file(self):
+        # small stable resource
+        url = self.http_client.server
+
+        path = self.http_client.download(
+            url,
+            file_name="index.html",
+        )
+
+        assert path.exists()
+        assert path.is_file()
+        assert path.read_text(encoding="utf-8")
+
+    @pytest.mark.uses_test_server()
+    def test_download_is_cached_on_disk(self):
+        url = self.http_client.server
+
+        path1 = self.http_client.download(
+            url,
+            file_name="cached.html",
+        )
+        mtime1 = path1.stat().st_mtime
+
+        # second call should NOT re-download
+        path2 = self.http_client.download(
+            url,
+            file_name="cached.html",
+        )
+        mtime2 = path2.stat().st_mtime
+
+        assert path1 == path2
+        assert mtime1 == mtime2
+
+    @pytest.mark.uses_test_server()
+    def test_download_respects_custom_handler(self):
+        url = self.http_client.server
+
+        def handler(response, path: Path, encoding: str):
+            path.write_text("HANDLED", encoding=encoding)
+            return path
+
+        path = self.http_client.download(
+            url,
+            handler=handler,
+            file_name="handled.txt",
+        )
+
+        assert path.exists()
+        assert path.read_text() == "HANDLED"