get_root|child|item_links should be json-like

* Allow media_type on get_links and get_single_link to
  an iterable
* Use `application/json", "application/geo+json" or None
  as the media_type for `get_root_link`, `get_child_links`
  and `get_item_links`

Author: jsignell

pystac/catalog.py

@@ -464,7 +464,10 @@ def get_child_links(self) -> list[Link]:
         Return:
             List[Link]: List of links of this catalog with ``rel == 'child'``
         """
-        return self.get_links(pystac.RelType.CHILD)
+        return self.get_links(
+            rel=pystac.RelType.CHILD,
+            media_type=[None, pystac.MediaType.GEOJSON, pystac.MediaType.JSON],
+        )
 
     def clear_children(self) -> None:
         """Removes all children from this catalog.
@@ -624,7 +627,10 @@ def get_item_links(self) -> list[Link]:
         Return:
             List[Link]: List of links of this catalog with ``rel == 'item'``
         """
-        return self.get_links(pystac.RelType.ITEM)
+        return self.get_links(
+            rel=pystac.RelType.ITEM,
+            media_type=[None, pystac.MediaType.GEOJSON, pystac.MediaType.JSON],
+        )
 
     def to_dict(
         self, include_self_link: bool = True, transform_hrefs: bool = True

pystac/stac_object.py

@@ -3,12 +3,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable
 from html import escape
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    TypeVar,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, cast
 
 import pystac
 from pystac import STACError
@@ -27,6 +22,8 @@
 
 S = TypeVar("S", bound="STACObject")
 
+STACObject_MediaType: TypeAlias = str | pystac.MediaType | None
+
 
 class STACObjectType(StringEnum):
     CATALOG = "Catalog"
@@ -177,7 +174,7 @@ def traverse(obj: str | STACObject, visited: set[str | STACObject]) -> bool:
     def get_single_link(
         self,
         rel: str | pystac.RelType | None = None,
-        media_type: str | pystac.MediaType | None = None,
+        media_type: STACObject_MediaType | Iterable[STACObject_MediaType] = None,
     ) -> Link | None:
         """Get a single :class:`~pystac.Link` instance associated with this
         object.
@@ -186,37 +183,41 @@ def get_single_link(
             rel : If set, filter links such that only those
                 matching this relationship are returned.
             media_type: If set, filter the links such that only
-                those matching media_type are returned
+                those matching media_type are returned. media_type can
+                be a single value or a list of values.
 
         Returns:
-            Optional[:class:`~pystac.Link`]: First link that matches ``rel``
+            :class:`~pystac.Link` | None: First link that matches ``rel``
             and/or ``media_type``, or else the first link associated with
             this object.
         """
         if rel is None and media_type is None:
             return next(iter(self.links), None)
+        if media_type and isinstance(media_type, (str, pystac.MediaType)):
+            media_type = [media_type]
         return next(
             (
                 link
                 for link in self.links
                 if (rel is None or link.rel == rel)
-                and (media_type is None or link.media_type == media_type)
+                and (media_type is None or link.media_type in media_type)
             ),
             None,
         )
 
     def get_links(
         self,
         rel: str | pystac.RelType | None = None,
-        media_type: str | pystac.MediaType | None = None,
+        media_type: STACObject_MediaType | Iterable[STACObject_MediaType] = None,
     ) -> list[Link]:
         """Gets the :class:`~pystac.Link` instances associated with this object.
 
         Args:
             rel : If set, filter links such that only those
                 matching this relationship are returned.
             media_type: If set, filter the links such that only
-                those matching media_type are returned
+                those matching media_type are returned. media_type can
+                be a single value or a list of values.
 
         Returns:
             List[:class:`~pystac.Link`]: A list of links that match ``rel`` and/
@@ -225,13 +226,14 @@ def get_links(
         """
         if rel is None and media_type is None:
             return self.links
-        else:
-            return [
-                link
-                for link in self.links
-                if (rel is None or link.rel == rel)
-                and (media_type is None or link.media_type == media_type)
-            ]
+        if media_type and isinstance(media_type, (str, pystac.MediaType)):
+            media_type = [media_type]
+        return [
+            link
+            for link in self.links
+            if (rel is None or link.rel == rel)
+            and (media_type is None or link.media_type in media_type)
+        ]
 
     def clear_links(self, rel: str | pystac.RelType | None = None) -> None:
         """Clears all :class:`~pystac.Link` instances associated with this object.
@@ -252,7 +254,10 @@ def get_root_link(self) -> Link | None:
             :class:`~pystac.Link` or None: The root link for this object,
             or ``None`` if no root link is set.
         """
-        return self.get_single_link(pystac.RelType.ROOT)
+        return self.get_single_link(
+            rel=pystac.RelType.ROOT,
+            media_type=[None, pystac.MediaType.GEOJSON, pystac.MediaType.JSON],
+        )
 
     @property
     def self_href(self) -> str:

tests/test_catalog.py

@@ -1407,6 +1407,10 @@ def test_get_links(self) -> None:
             len(catalog.get_links(rel="search", media_type="application/geo+json")) == 1
         )
         assert len(catalog.get_links(media_type="text/html")) == 1
+        assert (
+            len(catalog.get_links(media_type=["text/html", "application/geo+json"]))
+            == 2
+        )
         assert len(catalog.get_links(rel="search")) == 2
         assert len(catalog.get_links(rel="via")) == 0
         assert len(catalog.get_links()) == 6
@@ -1968,3 +1972,47 @@ def test_add_item_layout_strategy(
     template = template.format(id=item.id).replace("$", "")
 
     assert item.self_href == f"{base_url}/{template}/{item_id}.json"
+
+
+def test_get_child_links_cares_about_media_type(catalog: pystac.Catalog) -> None:
+    catalog.links.extend(
+        [
+            pystac.Link(
+                rel="child", target="./child-1.json", media_type="application/json"
+            ),
+            pystac.Link(
+                rel="child", target="./child-2.json", media_type="application/geo+json"
+            ),
+            pystac.Link(rel="child", target="./child-3.json"),
+            # this one won't get counted since it's the wrong media_type
+            pystac.Link(rel="child", target="./child.html", media_type="text/html"),
+        ]
+    )
+
+    assert len(catalog.get_child_links()) == 3
+
+
+def test_get_item_links_cares_about_media_type(catalog: pystac.Catalog) -> None:
+    catalog.links.extend(
+        [
+            pystac.Link(
+                rel="item", target="./item-1.json", media_type="application/json"
+            ),
+            pystac.Link(
+                rel="item", target="./item-2.json", media_type="application/geo+json"
+            ),
+            pystac.Link(rel="item", target="./item-3.json"),
+            # this one won't get counted since it's the wrong media_type
+            pystac.Link(rel="item", target="./item.html", media_type="text/html"),
+        ]
+    )
+
+    assert len(catalog.get_item_links()) == 3
+
+
+def test_get_root_link_cares_about_media_type(catalog: pystac.Catalog) -> None:
+    catalog.links.insert(
+        0, pystac.Link(rel="root", target="./self.json", media_type="text/html")
+    )
+    root_link = catalog.get_root_link()
+    assert root_link and root_link.target != "./self.json"