_JobResultDownloader: support "auto_increment" var in path templates

CHANGELOG.md

@@ -12,7 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `AGENTS.md` with guidance for AI coding agents contributing to this project, including a request to disclose AI assistance in PRs ([#923](https://github.com/Open-EO/openeo-python-client/issues/923))
 - Add a `py.typed` to indicate to type checkers that the package contains type annotations.
 - Support document based "derived_from" links in `openeo.testing.results` ([#928](https://github.com/Open-EO/openeo-python-client/issues/928))
-- Add `JobResults.download_as_collection()` (experimental) to download job results as a self-contained STAC collection with rewritten hrefs ([#931](https://github.com/Open-EO/openeo-python-client/issues/931))
+- Add `JobResults.download_as_collection()` (experimental) to download job results as a self-contained STAC collection with rewritten hrefs, including a template system to fine-tune file names and paths ([#931](https://github.com/Open-EO/openeo-python-client/issues/931))
 
 ### Changed
 

openeo/_version.py

@@ -1 +1 @@
-__version__ = "0.52.0a4"
+__version__ = "0.52.0a5"

openeo/rest/_testing.py

@@ -537,7 +537,7 @@ def setup_job_results(
                 collection_assets[f"{item_id}-{asset_key}"] = asset
 
             item_href = self.setup_item(job_id=job_id, item_id=item_id, item_data=item_data, assets=assets)
-            collection_links.append({"rel": "item", "href": item_href})
+            collection_links.append({"rel": "item", "href": item_href, **item_data.get("link_extra", {})})
 
         for doc in linked_docs:
             collection_links.append(self.setup_linked_document(job_id=job_id, doc=doc))
@@ -586,8 +586,7 @@ def setup_asset(self, *, job_id: str, asset_data: dict) -> dict:
             self.requests_mock.head(href, headers={"Content-Length": f"{len(content)}"})
             self.requests_mock.get(href, content=content)
         return StacDummyBuilder.asset(
-            href=href,
-            type=asset_data.get("type", "image/tiff; application=geotiff"),
+            href=href, type=asset_data.get("type", "image/tiff; application=geotiff"), **asset_data.get("extra", {})
         )
 
     def setup_linked_document(self, *, job_id: str, doc: dict):
@@ -599,4 +598,4 @@ def setup_linked_document(self, *, job_id: str, doc: dict):
             text = doc.get("text", "hello world")
         self.requests_mock.head(href, headers={"Content-Length": f"{len(text)}"})
         self.requests_mock.get(href, text=text)
-        return {"rel": doc.get("rel", "doc"), "href": href}
+        return {"rel": doc.get("rel", "doc"), "href": href, **doc.get("link_extra", {})}

openeo/rest/job.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import collections
 import contextlib
 import copy
 import datetime
@@ -32,6 +33,7 @@
 from openeo.rest.models.general import LogsResponse
 from openeo.rest.models.logs import log_level_name
 from openeo.util import ensure_dir, ensure_parent_dir_for
+from openeo.utils.datastructure import make_new_key
 from openeo.utils.events import EVENTS
 from openeo.utils.http import (
     HTTP_502_BAD_GATEWAY,
@@ -418,6 +420,11 @@ def _filename_from_url(url: str, *, full: bool = False) -> str:
     return "/".join(parts)
 
 
+def _filename_extension_from_url(url: str, *, fallback: str = "") -> str:
+    filename = _filename_from_url(url)
+    return "".join(Path(filename).suffixes) or fallback
+
+
 _MEDIA_TYPE_EXTENSION_MAP = {
     "image/tiff": ".tiff",
     "image/tiff; application=geotiff": ".tiff",
@@ -685,9 +692,10 @@ def download_as_collection(
         target: Union[Path, str, None] = None,
         *,
         rewrite_references: bool = True,
+        add_original_hrefs: bool = False,
         download_derived_from: bool = False,
         download_collection_assets: bool = False,
-        json_dumping: Optional[dict] = None,
+        json_dump: Optional[dict] = None,
         on_download_failure: Literal["warn", "raise"] = "warn",
         path_templates: Optional[dict] = None,
         redact_url_logging: bool = True,
@@ -705,13 +713,14 @@ def download_as_collection(
         :param rewrite_references: whether to rewrite (item/asset/...) HREFs
             in the downloaded STAC documents to point to the local files
             instead of the original URLs.
+        :param add_original_hrefs: whether to add the original HREF URLs as "alternate" HREFs in the metadata
         :param download_derived_from: whether to download
             additional "derived_from" documents linked from the STAC collection.
         :param download_collection_assets: whether to download
             the STAC Collection level assets in addition to assets from linked STAC Items.
-        :param json_dumping: kwargs to finetune json.dump when writing STAC metadata files.
+        :param json_dump: kwargs to finetune json.dump when writing STAC metadata files.
         :param on_download_failure: how to handle download failures, one of "warn" or "raise".
-        :param path_templates: optional template overrides for download paths.
+        :param path_templates: optional dictionary of template overrides for download paths.
         :param redact_url_logging: whether to redact (possibly sensitive) query parameters from URLs in logging,
 
         .. versionadded:: 0.52.0
@@ -720,7 +729,8 @@ def download_as_collection(
             job=self._job,
             target=target,
             rewrite_references=rewrite_references,
-            json_dumping=json_dumping,
+            add_original_hrefs=add_original_hrefs,
+            json_dump=json_dump,
             on_download_failure=on_download_failure,
             path_templates=path_templates,
             redact_url_logging=redact_url_logging,
@@ -758,8 +768,10 @@ def register(self, path: Path):
 
 class _JobResultDownloader:
     """
-    Helper class to download batch job results as a STAC collection (openEO API 1.1 style):
-    recursively walking through items, assets and additional linked metadata.
+    Helper class to download batch job results
+    as a local, self-contained STAC collection (openEO API 1.1 style):
+    recursively walking through items, assets, additional linked metadata, etc.,
+    and rewriting the related links accordingly.
 
     .. warning:: this is an experimental API, subject to change.
 
@@ -788,7 +800,8 @@ def __init__(
         job: BatchJob,
         target: Union[Path, str, None] = None,
         rewrite_references: bool = True,
-        json_dumping: Optional[dict] = None,
+        add_original_hrefs: bool = False,
+        json_dump: Optional[dict] = None,
         on_download_failure: Literal["warn", "raise"] = "warn",
         path_templates: Optional[dict] = None,
         redact_url_logging: bool = True,
@@ -799,18 +812,20 @@ def __init__(
         if self._root_dir.exists() and not self._root_dir.is_dir():
             raise OpenEoClientException(f"Download target {self._root_dir} exists but isn't a folder.")
         self._rewrite_references = rewrite_references
+        self._add_original_hrefs = add_original_hrefs
         # TODO: also support passing a `json.dump`-style callable to customize json dumping
-        self._json_dumping = {"ensure_ascii": False, **(json_dumping or {})}
+        self._json_dump = {"ensure_ascii": False, **(json_dump or {})}
         self._download_tracker = _DownloadTracker()
         self._on_download_failure = on_download_failure
         self._path_templates = {**self.DEFAULT_PATH_TEMPLATES, **(path_templates or {})}
         self._redact = get_url_query_param_stripper(redact_url_logging)
+        self._auto_increment_counters: Dict[str, int] = collections.defaultdict(int)
 
     def _write_json_file(self, data: dict, path: Union[str, Path]) -> Path:
         path = Path(path)
         ensure_parent_dir_for(path)
         with open(path, mode="w", encoding="utf-8") as f:
-            json.dump(obj=data, fp=f, **self._json_dumping)
+            json.dump(obj=data, fp=f, **self._json_dump)
         return path
 
     @contextlib.contextmanager
@@ -833,6 +848,8 @@ def build_path_collection(self, *, collection_id: str) -> Path:
         vars = {
             "job_id": _sanitize_filename(self._job.job_id),
             "collection_id": _sanitize_filename(collection_id),
+            "auto_increment": self._auto_increment_id("collection"),
+            "extension": ".json",
         }
         return self._root_dir / self._path_templates["collection"].format(**vars)
 
@@ -841,6 +858,8 @@ def build_path_item(self, *, item_id: str) -> Path:
         vars = {
             "job_id": _sanitize_filename(self._job.job_id),
             "item_id": _sanitize_filename(item_id),
+            "auto_increment": self._auto_increment_id("item"),
+            "extension": ".json",
         }
         return self._root_dir / self._path_templates["item"].format(**vars)
 
@@ -850,6 +869,10 @@ def build_path_asset(self, *, asset_key: str, asset_href: str, item_id: Optional
             "job_id": _sanitize_filename(self._job.job_id),
             "asset_key": _sanitize_filename(asset_key),
             "asset_filename": _filename_from_url(asset_href, full=False),
+            # TODO: separate pool for item- and collection-assets?
+            "auto_increment": self._auto_increment_id("asset"),
+            # TODO: also leverage media type to determine extension?
+            "extension": _filename_extension_from_url(asset_href),
         }
         if item_id:
             vars["item_id"] = _sanitize_filename(item_id)
@@ -862,9 +885,17 @@ def build_path_generic_link(self, *, rel: str, href: str) -> Path:
             "job_id": _sanitize_filename(self._job.job_id),
             "rel": _sanitize_filename(rel),
             "filename": _filename_from_url(href, full=False),
+            "auto_increment": self._auto_increment_id("generic"),
+            # TODO: also leverage media type to determine extension?
+            "extension": _filename_extension_from_url(href),
         }
         return self._root_dir / self._path_templates["generic-link"].format(**vars)
 
+    def _auto_increment_id(self, pool: str) -> int:
+        """Generate auto-incrementing ID within a given pool of entity types."""
+        self._auto_increment_counters[pool] += 1
+        return self._auto_increment_counters[pool]
+
     def _relative_to(self, target: Path, doc: Path) -> str:
         """Get relative reference to target to be used from given document"""
         return Path(os.path.relpath(target, start=doc.parent)).as_posix()
@@ -898,10 +929,7 @@ def download_collection(
                 with self._download_attempt_context(name=f"item {link=}"):
                     path = self._download_item(href=link["href"])
                     if self._rewrite_references:
-                        rel_path = self._relative_to(target=path, doc=result_metadata_path)
-                        logger.debug(f"Rewriting link {self._redact(link)=} href to local {rel_path=}")
-                        link["href"] = rel_path
-
+                        self._rewrite_href(obj=link, path=path, relative_to=result_metadata_path)
             elif link["rel"] in extra_rels:
                 with self._download_attempt_context(name=f"link {link=}"):
                     path = self.build_path_generic_link(rel=link["rel"], href=link["href"])
@@ -910,9 +938,7 @@ def download_collection(
                     logger.debug(f"Downloaded link {self._redact(link)=} to {path=}")
                     self._download_tracker.register(path)
                     if self._rewrite_references:
-                        rel_path = self._relative_to(target=path, doc=result_metadata_path)
-                        logger.debug(f"Rewriting link {self._redact(link)=} href to local {rel_path=}")
-                        link["href"] = rel_path
+                        self._rewrite_href(obj=link, path=path, relative_to=result_metadata_path)
 
         if download_collection_assets:
             for asset_key, asset in result_metadata.get("assets", {}).items():
@@ -921,9 +947,7 @@ def download_collection(
                         asset_key=asset_key, asset_href=asset["href"], asset_metadata=asset, item_id=None
                     )
                     if self._rewrite_references:
-                        rel_path = self._relative_to(target=path, doc=result_metadata_path)
-                        logger.debug(f"Rewriting STAC Collection asset {asset_key=} href to local {rel_path=}")
-                        asset["href"] = rel_path
+                        self._rewrite_href(obj=asset, path=path, relative_to=result_metadata_path)
 
         if self._rewrite_references:
             # Rewrite the root collection metadata with updated references
@@ -934,6 +958,16 @@ def download_collection(
 
         return self._download_tracker.paths
 
+    def _rewrite_href(self, obj: dict, path: Path, relative_to: Path):
+        """Rewrite (in-place) href of given object."""
+        original = obj["href"]
+        rel_path = self._relative_to(target=path, doc=relative_to)
+        logger.debug(f"Rewriting href {rel_path=} (in {self._redact(obj)})")
+        obj["href"] = rel_path
+        if self._add_original_hrefs:
+            alternate = obj.setdefault("alternate", {})
+            alternate[make_new_key(alternate, "original")] = {"href": original}
+
     def _download_item(self, href: str) -> Path:
         item: dict = self._connection.get(href, expected_status=200).json()
         item_id = item["id"]
@@ -948,9 +982,7 @@ def _download_item(self, href: str) -> Path:
                     asset_key=asset_key, asset_href=asset["href"], asset_metadata=asset, item_id=item_id
                 )
                 if self._rewrite_references:
-                    rel_path = self._relative_to(target=asset_path, doc=metadata_path)
-                    logger.debug(f"Rewriting asset {asset_key=} ({item_id=}) href to local {rel_path=}")
-                    asset["href"] = rel_path
+                    self._rewrite_href(obj=asset, path=asset_path, relative_to=metadata_path)
 
         if self._rewrite_references:
             logger.info(f"Update write of STAC Item {item_id!r} metadata to {metadata_path}")

openeo/utils/datastructure.py

@@ -0,0 +1,18 @@
+"""
+Generic data structure related utilities
+"""
+
+import itertools
+from typing import Container
+
+
+def make_new_key(container: Container[str], key: str, *, start: int = 1):
+    """
+    Construct a new key that doesn't exist yet in the given container (dict, set, ...),
+    based on given base key (to use directly if possible, or by appending an auto-increment-style suffix).
+    """
+    if key not in container:
+        return key
+    for i in itertools.count(start):
+        if (k := f"{key}-{i}") not in container:
+            return k