Add a post-init modifier keyword

CHANGELOG.md

@@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 ## [Unreleased] - TBD
 
+### Added
+
+- Added a new keyword `modifier` to various constructors like `Client.open()` [#259](https://github.com/stac-utils/pystac-client/issues/259)
+
 ### Fixed
 
 - Fix type annotation of `Client._stac_io` and avoid implicit re-exports in `pystac_client.__init__.py` [#249](https://github.com/stac-utils/pystac-client/pull/249)

docs/usage.rst

@@ -288,3 +288,33 @@ descending sort and a ``+`` prefix or no prefix means an ascending sort.
                 {"direction": "asc", "field": "collection"},
             ]
     ... )
+
+Automatically modifying results
+-------------------------------
+
+Some systems, like the `Microsoft Planetary Computer <http://planetarycomputer.microsoft.com/>`__,
+have public STAC metadata but require require some `authentication <https://planetarycomputer.microsoft.com/docs/concepts/sas/>`__
+to access the actual assets.
+
+``pystac-client`` provides a ``modifier`` keyword that can automatically
+modify the STAC objects returned by the STAC API.
+
+.. code-block:: python
+
+   >>> from pystac_client import Client
+   >>> import planetary_computer, requests
+   >>> api = Client.open(
+   ...    'https://planetarycomputer.microsoft.com/api/stac/v1',
+   ...    modifier=planetary_computer.sign_inplace,
+   ... )
+   >>> item = next(catalog.get_collection("sentinel-2-l2a").get_all_items())
+   >>> requests.head(item.assets["B02"].href).status_code
+   200
+
+Without the modifier, we would have received a 404 error because the asset
+is in a private storage container.
+
+``pystac-client`` expects that the ``modifier`` callable modifies the result
+object in-place and returns no result. A warning is emitted if your
+``modifier`` returns a non-None result that is not the same object as the
+input.

pystac_client/_utils.py

@@ -0,0 +1,25 @@
+import warnings
+from typing import Callable, Optional, Union
+
+import pystac
+
+from pystac_client.errors import IgnoredResultWarning
+
+Modifiable = Union[pystac.Collection, pystac.Item, pystac.ItemCollection, dict]
+
+
+def call_modifier(
+    modifier: Optional[Callable[[Modifiable], None]], obj: Modifiable
+) -> None:
+    """Calls the user's modifier and validates that the result is None."""
+    if modifier is None:
+        return None
+
+    result = modifier(obj)
+    if result is not None and result is not obj:
+        warnings.warn(
+            f"modifier '{modifier}' returned a result that's being ignored. "
+            "You should ensure that 'modifier' is operating in-place and use "
+            "a function that returns 'None' or silence this warning.",
+            IgnoredResultWarning,
+        )

pystac_client/client.py

@@ -1,10 +1,21 @@
 from functools import lru_cache
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Union,
+    cast,
+)
 
 import pystac
 import pystac.validation
-from pystac import Collection
+from pystac import CatalogType, Collection
 
+from pystac_client._utils import Modifiable, call_modifier
 from pystac_client.collection_client import CollectionClient
 from pystac_client.conformance import ConformanceClasses
 from pystac_client.errors import ClientTypeError
@@ -46,6 +57,31 @@ class Client(pystac.Catalog):
 
     _stac_io: Optional[StacApiIO]
 
+    def __init__(
+        self,
+        id: str,
+        description: str,
+        title: Optional[str] = None,
+        stac_extensions: Optional[List[str]] = None,
+        extra_fields: Optional[Dict[str, Any]] = None,
+        href: Optional[str] = None,
+        catalog_type: CatalogType = CatalogType.ABSOLUTE_PUBLISHED,
+        *,
+        modifier: Optional[Callable[[Modifiable], None]] = None,
+        **kwargs: Dict[str, Any],
+    ):
+        super().__init__(
+            id,
+            description,
+            title=title,
+            stac_extensions=stac_extensions,
+            extra_fields=extra_fields,
+            href=href,
+            catalog_type=catalog_type,
+            **kwargs,
+        )
+        self.modifier = modifier
+
     def __repr__(self) -> str:
         return "<Client id={}>".format(self.id)
 
@@ -56,6 +92,7 @@ def open(
         headers: Optional[Dict[str, str]] = None,
         parameters: Optional[Dict[str, Any]] = None,
         ignore_conformance: bool = False,
+        modifier: Optional[Callable[[Modifiable], None]] = None,
     ) -> "Client":
         """Opens a STAC Catalog or API
         This function will read the root catalog of a STAC Catalog or API
@@ -72,11 +109,33 @@ def open(
                 functions will skip checking conformance, and may throw an unknown
                 error if that feature is
                 not supported, rather than a :class:`NotImplementedError`.
+            modifier : A callable that modifies the children collection and items
+                returned by this Client. This can be useful for injecting
+                authentication parameters into child assets to access data
+                from non-public sources.
+
+                The callable should expect a single argument, which will be one
+                of the following types:
+
+                * :class:`pystac.Collection`
+                * :class:`pystac.Item`
+                * :class:`pystac.ItemCollection`
+                * A STAC item-like :class:`dict`
+                * A STAC collection-like :class:`dict`
+
+                The callable should mutate the argument in place and return ``None``.
+
+                ``modifier`` propagates recursively to children of this Client.
+                After getting a child collection with, e.g.
+                :meth:`Client.get_collection`, the child items of that collection
+                will still be signed with ``modifier``.
 
         Return:
             catalog : A :class:`Client` instance for this Catalog/API
         """
-        client: Client = cls.from_file(url, headers=headers, parameters=parameters)
+        client: Client = cls.from_file(
+            url, headers=headers, parameters=parameters, modifier=modifier
+        )
         search_link = client.get_search_link()
         # if there is a search link, but no conformsTo advertised, ignore
         # conformance entirely
@@ -102,6 +161,7 @@ def from_file(  # type: ignore
         stac_io: Optional[StacApiIO] = None,
         headers: Optional[Dict[str, str]] = None,
         parameters: Optional[Dict[str, Any]] = None,
+        modifier: Optional[Callable[[Modifiable], None]] = None,
     ) -> "Client":
         """Open a STAC Catalog/API
 
@@ -116,6 +176,7 @@ def from_file(  # type: ignore
         client._stac_io._conformance = client.extra_fields.get(  # type: ignore
             "conformsTo", []
         )
+        client.modifier = modifier
 
         return client
 
@@ -135,17 +196,25 @@ def from_dict(
         root: Optional[pystac.Catalog] = None,
         migrate: bool = False,
         preserve_dict: bool = True,
+        modifier: Optional[Callable[[Modifiable], None]] = None,
     ) -> "Client":
         try:
             # this will return a Client because we have used a StacApiIO instance
-            return super().from_dict(  # type: ignore
+            result = super().from_dict(
                 d=d, href=href, root=root, migrate=migrate, preserve_dict=preserve_dict
             )
         except pystac.STACTypeError:
             raise ClientTypeError(
                 f"Could not open Client (href={href}), "
                 f"expected type=Catalog, found type={d.get('type', None)}"
             )
+        # cast require for mypy to believe that we have a Client, rather than
+        # the super type.
+        # https://github.com/stac-utils/pystac/issues/862
+        result = cast(Client, result)
+
+        result.modifier = modifier
+        return result
 
     @lru_cache()
     def get_collection(self, collection_id: str) -> Optional[Collection]:
@@ -160,12 +229,16 @@ def get_collection(self, collection_id: str) -> Optional[Collection]:
         if self._supports_collections() and self._stac_io:
             url = f"{self.get_self_href()}/collections/{collection_id}"
             collection = CollectionClient.from_dict(
-                self._stac_io.read_json(url), root=self
+                self._stac_io.read_json(url),
+                root=self,
+                modifier=self.modifier,
             )
+            call_modifier(self.modifier, collection)
             return collection
         else:
             for col in self.get_collections():
                 if col.id == collection_id:
+                    call_modifier(self.modifier, col)
                     return col
 
         return None
@@ -179,15 +252,23 @@ def get_collections(self) -> Iterator[Collection]:
         Return:
             Iterator[Collection]: Iterator over Collections in Catalog/API
         """
+        collection: Union[Collection, CollectionClient]
+
         if self._supports_collections() and self.get_self_href() is not None:
             url = f"{self.get_self_href()}/collections"
             for page in self._stac_io.get_pages(url):  # type: ignore
                 if "collections" not in page:
                     raise APIError("Invalid response from /collections")
                 for col in page["collections"]:
-                    yield CollectionClient.from_dict(col, root=self)
+                    collection = CollectionClient.from_dict(
+                        col, root=self, modifier=self.modifier
+                    )
+                    call_modifier(self.modifier, collection)
+                    yield collection
         else:
-            yield from super().get_collections()
+            for collection in super().get_collections():
+                call_modifier(self.modifier, collection)
+                yield collection
 
     def get_items(self) -> Iterator["Item_Type"]:
         """Return all items of this catalog.
@@ -200,7 +281,9 @@ def get_items(self) -> Iterator["Item_Type"]:
             search = self.search()
             yield from search.items()
         else:
-            yield from super().get_items()
+            for item in super().get_items():
+                call_modifier(self.modifier, item)
+                yield item
 
     def get_all_items(self) -> Iterator["Item_Type"]:
         """Get all items from this catalog and all subcatalogs. Will traverse
@@ -212,9 +295,12 @@ def get_all_items(self) -> Iterator["Item_Type"]:
                 child links.
         """
         if self._conforms_to(ConformanceClasses.ITEM_SEARCH):
+            # these are already modified
             yield from self.get_items()
         else:
-            yield from super().get_items()
+            for item in super().get_items():
+                call_modifier(self.modifier, item)
+                yield item
 
     def search(
         self,
@@ -369,6 +455,7 @@ def search(
             filter_lang=filter_lang,
             sortby=sortby,
             fields=fields,
+            modifier=self.modifier,
         )
 
     def get_search_link(self) -> Optional[pystac.Link]: