Add type annotations to everything

Author: lossyrob

pystac/__init__.py

@@ -13,23 +13,25 @@ class STACError(Exception):
     pass
 
 
-from pystac.version import (__version__, get_stac_version, set_stac_version)
+from typing import Any, Dict, Optional
+from pystac.version import (__version__, get_stac_version, set_stac_version)  # type:ignore
 from pystac.stac_io import STAC_IO
-from pystac.extensions import Extensions
-from pystac.stac_object import (STACObject, STACObjectType)
-from pystac.media_type import MediaType
-from pystac.link import (Link, HIERARCHICAL_LINKS)
-from pystac.catalog import (Catalog, CatalogType)
-from pystac.collection import (Collection, Extent, SpatialExtent, TemporalExtent, Provider)
-from pystac.item import (Item, Asset, CommonMetadata)
+from pystac.extensions import Extensions  # type:ignore
+from pystac.stac_object import (STACObject, STACObjectType)  # type:ignore
+from pystac.media_type import MediaType  # type:ignore
+from pystac.link import (Link, HIERARCHICAL_LINKS)  # type:ignore
+from pystac.catalog import (Catalog, CatalogType)  # type:ignore
+from pystac.collection import (Collection, Extent, SpatialExtent, TemporalExtent, # type:ignore
+                               Provider)  # type:ignore
+from pystac.item import (Item, Asset, CommonMetadata)  # type:ignore
 
 from pystac.serialization import stac_object_from_dict
 
 import pystac.validation
 
 STAC_IO.stac_object_from_dict = stac_object_from_dict
 
-from pystac import extensions
+import pystac.extensions.base
 import pystac.extensions.eo
 import pystac.extensions.label
 import pystac.extensions.pointcloud
@@ -43,19 +45,24 @@ class STACError(Exception):
 import pystac.extensions.view
 import pystac.extensions.file
 
-STAC_EXTENSIONS = extensions.base.RegisteredSTACExtensions([
-    extensions.eo.EO_EXTENSION_DEFINITION, extensions.label.LABEL_EXTENSION_DEFINITION,
-    extensions.pointcloud.POINTCLOUD_EXTENSION_DEFINITION,
-    extensions.projection.PROJECTION_EXTENSION_DEFINITION, extensions.sar.SAR_EXTENSION_DEFINITION,
-    extensions.sat.SAT_EXTENSION_DEFINITION, extensions.scientific.SCIENTIFIC_EXTENSION_DEFINITION,
-    extensions.single_file_stac.SFS_EXTENSION_DEFINITION,
-    extensions.timestamps.TIMESTAMPS_EXTENSION_DEFINITION,
-    extensions.version.VERSION_EXTENSION_DEFINITION, extensions.view.VIEW_EXTENSION_DEFINITION,
-    extensions.file.FILE_EXTENSION_DEFINITION
-])
-
-
-def read_file(href):
+STAC_EXTENSIONS: pystac.extensions.base.RegisteredSTACExtensions = pystac.extensions.base.RegisteredSTACExtensions(
+    [
+        pystac.extensions.eo.EO_EXTENSION_DEFINITION,
+        pystac.extensions.label.LABEL_EXTENSION_DEFINITION,
+        pystac.extensions.pointcloud.POINTCLOUD_EXTENSION_DEFINITION,
+        pystac.extensions.projection.PROJECTION_EXTENSION_DEFINITION,
+        pystac.extensions.sar.SAR_EXTENSION_DEFINITION,
+        pystac.extensions.sat.SAT_EXTENSION_DEFINITION,
+        pystac.extensions.scientific.SCIENTIFIC_EXTENSION_DEFINITION,
+        pystac.extensions.single_file_stac.SFS_EXTENSION_DEFINITION,
+        pystac.extensions.timestamps.TIMESTAMPS_EXTENSION_DEFINITION,
+        pystac.extensions.version.VERSION_EXTENSION_DEFINITION,
+        pystac.extensions.view.VIEW_EXTENSION_DEFINITION,
+        pystac.extensions.file.FILE_EXTENSION_DEFINITION
+    ])
+
+
+def read_file(href: str) -> STACObject:
     """Reads a STAC object from a file.
 
     This method will return either a Catalog, a Collection, or an Item based on what the
@@ -73,7 +80,7 @@ def read_file(href):
     return STACObject.from_file(href)
 
 
-def write_file(obj, include_self_link=True, dest_href=None):
+def write_file(obj: STACObject, include_self_link: bool = True, dest_href: Optional[str] = None):
     """Writes a STACObject to a file.
 
     This will write only the Catalog, Collection or Item ``obj``. It will not attempt
@@ -96,7 +103,7 @@ def write_file(obj, include_self_link=True, dest_href=None):
     obj.save_object(include_self_link=include_self_link, dest_href=dest_href)
 
 
-def read_dict(d, href=None, root=None):
+def read_dict(d: Dict[str, Any], href: Optional[str] = None, root: Optional[Catalog] = None):
     """Reads a STAC object from a dict representing the serialized JSON version of the
     STAC object.
 

pystac/cache.py

@@ -1,10 +1,13 @@
 from collections import ChainMap
 from copy import copy
+from pystac.collection import Collection
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from pystac.stac_object import STACObject
 
 import pystac
 
 
-def get_cache_key(stac_object):
+def get_cache_key(stac_object: STACObject) -> Tuple[str, bool]:
     """Produce a cache key for the given STAC object.
 
     If a self href is set, use that as the cache key.
@@ -20,7 +23,7 @@ def get_cache_key(stac_object):
     if href is not None:
         return (href, True)
     else:
-        ids = []
+        ids: List[str] = []
         obj = stac_object
         while obj is not None:
             ids.append(obj.id)
@@ -52,14 +55,17 @@ class ResolvedObjectCache:
             their cached object.
         ids_to_collections (Dict[str, Collection]): Map of collection IDs to collections.
     """
-    def __init__(self, id_keys_to_objects=None, hrefs_to_objects=None, ids_to_collections=None):
+    def __init__(self,
+                 id_keys_to_objects: Optional[Dict[str, STACObject]] = None,
+                 hrefs_to_objects: Optional[Dict[str, STACObject]] = None,
+                 ids_to_collections: Dict[str, Collection] = None):
         self.id_keys_to_objects = id_keys_to_objects or {}
         self.hrefs_to_objects = hrefs_to_objects or {}
         self.ids_to_collections = ids_to_collections or {}
 
         self._collection_cache = None
 
-    def get_or_cache(self, obj):
+    def get_or_cache(self, obj: STACObject) -> STACObject:
         """Gets the STACObject that is the cached version of the given STACObject; or, if
         none exists, sets the cached object to the given object.
 
@@ -85,7 +91,7 @@ def get_or_cache(self, obj):
                 self.cache(obj)
                 return obj
 
-    def get(self, obj):
+    def get(self, obj: STACObject) -> Optional[STACObject]:
         """Get the cached object that has the same cache key as the given object.
 
         Args:
@@ -101,7 +107,7 @@ def get(self, obj):
         else:
             return self.id_keys_to_objects.get(key)
 
-    def get_by_href(self, href):
+    def get_by_href(self, href: str) -> Optional[STACObject]:
         """Gets the cached object at href.
 
         Args:
@@ -112,7 +118,7 @@ def get_by_href(self, href):
         """
         return self.hrefs_to_objects.get(href)
 
-    def get_collection_by_id(self, id):
+    def get_collection_by_id(self, id: str) -> Optional[Collection]:
         """Retrieved a cached Collection by its ID.
 
         Args:
@@ -124,7 +130,7 @@ def get_collection_by_id(self, id):
         """
         return self.ids_to_collections.get(id)
 
-    def cache(self, obj):
+    def cache(self, obj: STACObject) -> None:
         """Set the given object into the cache.
 
         Args:
@@ -136,10 +142,10 @@ def cache(self, obj):
         else:
             self.id_keys_to_objects[key] = obj
 
-        if obj.STAC_OBJECT_TYPE == pystac.STACObjectType.COLLECTION:
+        if isinstance(obj, Collection):
             self.ids_to_collections[obj.id] = obj
 
-    def remove(self, obj):
+    def remove(self, obj: STACObject) -> None:
         """Removes any cached object that matches the given object's cache key.
 
         Args:
@@ -155,21 +161,21 @@ def remove(self, obj):
         if obj.STAC_OBJECT_TYPE == pystac.STACObjectType.COLLECTION:
             self.id_keys_to_objects.pop(obj.id, None)
 
-    def __contains__(self, obj):
+    def __contains__(self, obj: STACObject) -> bool:
         key, is_href = get_cache_key(obj)
         return key in self.hrefs_to_objects if is_href else key in self.id_keys_to_objects
 
-    def contains_collection_id(self, collection_id):
+    def contains_collection_id(self, collection_id: str) -> bool:
         """Returns True if there is a collection with given collection ID is cached."""
         return collection_id in self.ids_to_collections
 
-    def as_collection_cache(self):
+    def as_collection_cache(self) -> "CollectionCache":
         if self._collection_cache is None:
             self._collection_cache = ResolvedObjectCollectionCache(self)
         return self._collection_cache
 
     @staticmethod
-    def merge(first, second):
+    def merge(first: "ResolvedObjectCache", second: "ResolvedObjectCache") -> "ResolvedObjectCache":
         """Merges two ResolvedObjectCache.
 
         The merged cache will give preference to the first argument; that is, if there
@@ -206,55 +212,65 @@ class CollectionCache:
     The CollectionCache will contain collections as either as dicts or PySTAC Collections,
     and will set Collection JSON that it reads in order to merge in common properties.
     """
-    def __init__(self, cached_ids=None, cached_hrefs=None):
+    def __init__(self,
+                 cached_ids: Dict[str, Union[Collection, Dict[str, Any]]] = None,
+                 cached_hrefs: Dict[str, Union[Collection, Dict[str, Any]]] = None):
         self.cached_ids = cached_ids or {}
         self.cached_hrefs = cached_hrefs or {}
 
-    def get_by_id(self, collection_id):
+    def get_by_id(self, collection_id: str) -> Optional[Union[Collection, Dict[str, Any]]]:
         return self.cached_ids.get(collection_id)
 
-    def get_by_href(self, href):
+    def get_by_href(self, href: str) -> Optional[Union[Collection, Dict[str, Any]]]:
         return self.cached_hrefs.get(href)
 
-    def contains_id(self, collection_id):
+    def contains_id(self, collection_id: str) -> bool:
         return collection_id in self.cached_ids
 
-    def cache(self, collection, href=None):
+    def cache(self, collection: Union[Collection, Dict[str, Any]], href: Optional[str] = None) -> None:
         """Caches a collection JSON."""
-        self.cached_ids[collection['id']] = collection
+        if isinstance(collection, Collection):
+            self.cached_ids[collection.id] = collection
+        else:
+            self.cached_ids[collection['id']] = collection
 
         if href is not None:
             self.cached_hrefs[href] = collection
 
 
 class ResolvedObjectCollectionCache(CollectionCache):
-    def __init__(self, resolved_object_cache, cached_ids=None, cached_hrefs=None):
+    def __init__(self,
+                 resolved_object_cache: ResolvedObjectCache,
+                 cached_ids: Dict[str, Union[Collection, Dict[str, Any]]] = None,
+                 cached_hrefs: Dict[str, Union[Collection, Dict[str, Any]]] = None):
         super().__init__(cached_ids, cached_hrefs)
         self.resolved_object_cache = resolved_object_cache
 
-    def get_by_id(self, collection_id):
+    def get_by_id(self, collection_id: str) -> Optional[Union[Collection, Dict[str, Any]]]:
         result = self.resolved_object_cache.get_collection_by_id(collection_id)
         if result is None:
             return super().get_by_id(collection_id)
         else:
             return result
 
-    def get_by_href(self, href):
+    def get_by_href(self, href: str) -> Optional[Union[Collection, Dict[str, Any]]]:
         result = self.resolved_object_cache.get_by_href(href)
         if result is None:
             return super().get_by_href(href)
         else:
-            return result
+            return cast(Collection, result)
 
-    def contains_id(self, collection_id):
+    def contains_id(self, collection_id: str) -> bool:
         return (self.resolved_object_cache.contains_collection_id(collection_id)
                 or super().contains_id(collection_id))
 
-    def cache(self, collection, href=None):
+    def cache(self, collection: Dict[str, Any], href: Optional[str] = None) -> None:
         super().cache(collection, href)
 
     @staticmethod
-    def merge(resolved_object_cache, first, second):
+    def merge(resolved_object_cache: ResolvedObjectCache,
+              first: Optional["ResolvedObjectCollectionCache"],
+              second: Optional["ResolvedObjectCollectionCache"]) -> "ResolvedObjectCollectionCache":
         first_cached_ids = {}
         if first is not None:
             first_cached_ids = copy(first.cached_ids)