Merge pull request #511 from duckontheweb/add/gh-169-get-collections

Add get_collections + get_all_collections methods to Catalog

Author: Jon Duckworth

CHANGELOG.md

@@ -14,6 +14,9 @@
 - Define equality and `__repr__` of `RangeSummary` instances based on `to_dict`
   representation ([#513](https://github.com/stac-utils/pystac/pull/513))
 - Sat Extension summaries ([#509](https://github.com/stac-utils/pystac/pull/509))
+- `Catalog.get_collections` for getting all child
+  `Collections` for a catalog, and `Catalog.get_all_collections` for recursively getting
+  all child `Collections` for a catalog and its children ([#511](https://github.com/stac-utils/pystac/pull/))
 
 ### Changed
 

pystac/catalog.py

@@ -320,6 +320,21 @@ def get_children(self) -> Iterable[Union["Catalog", "Collection_Type"]]:
             self.get_stac_objects(pystac.RelType.CHILD),
         )
 
+    def get_collections(self) -> Iterable["Collection_Type"]:
+        """Return all children of this catalog that are :class:`~pystac.Collection`
+        instances."""
+        return map(
+            lambda x: cast(pystac.Collection, x),
+            self.get_stac_objects(pystac.RelType.CHILD, pystac.Collection),
+        )
+
+    def get_all_collections(self) -> Iterable["Collection_Type"]:
+        """Get all collections from this catalog and all subcatalogs. Will traverse
+        any subcatalogs recursively."""
+        yield from self.get_collections()
+        for child in self.get_children():
+            yield from child.get_collections()
+
     def get_child_links(self) -> List[Link]:
         """Return all child links of this catalog.
 

pystac/stac_object.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Any, Dict, Iterable, List, Optional, cast, TYPE_CHECKING, Union
+from typing import Any, Dict, Iterable, List, Optional, Type, cast, TYPE_CHECKING, Union
 
 import pystac
 from pystac import STACError
@@ -272,25 +272,29 @@ def set_parent(self, parent: Optional["Catalog_Type"]) -> None:
             self.add_link(Link.parent(parent))
 
     def get_stac_objects(
-        self, rel: Union[str, pystac.RelType]
+        self, rel: Union[str, pystac.RelType], typ: Optional[Type["STACObject"]] = None
     ) -> Iterable["STACObject"]:
         """Gets the :class:`~pystac.STACObject` instances that are linked to
         by links with their ``rel`` property matching the passed in argument.
 
         Args:
             rel : The relation to match each :class:`~pystac.Link`'s
                 ``rel`` property against.
+            typ : If not ``None``, objects will only be yielded if they are instances of
+                ``typ``.
 
         Returns:
             Iterable[STACObjects]: A possibly empty iterable of STACObjects that are
-            connected to this object through links with the given ``rel``.
+            connected to this object through links with the given ``rel`` and are of
+            type ``typ`` (if given).
         """
         links = self.links[:]
         for i in range(0, len(links)):
             link = links[i]
             if link.rel == rel:
                 link.resolve_stac_object(root=self.get_root())
-                yield cast("STACObject", link.target)
+                if typ is None or isinstance(link.target, typ):
+                    yield cast("STACObject", link.target)
 
     def save_object(
         self,

tests/test_catalog.py

@@ -996,6 +996,20 @@ def testfrom_invalid_dict_raises_exception(self) -> None:
         with self.assertRaises(pystac.STACTypeError):
             _ = pystac.Catalog.from_dict(collection_dict)
 
+    def test_get_collections(self) -> None:
+        catalog = TestCases.test_case_2()
+        collections = list(catalog.get_collections())
+
+        self.assertGreater(len(collections), 0)
+        self.assertTrue(all(isinstance(c, pystac.Collection) for c in collections))
+
+    def test_get_all_collections(self) -> None:
+        catalog = TestCases.test_case_1()
+        all_collections = list(catalog.get_all_collections())
+
+        self.assertGreater(len(all_collections), 0)
+        self.assertTrue(all(isinstance(c, pystac.Collection) for c in all_collections))
+
 
 class FullCopyTest(unittest.TestCase):
     def check_link(self, link: pystac.Link, tag: str) -> None: