def remove_extension(self, extension_id):
"""Remove an extension from PySTAC."""
if extension_id not in self.extensions:
raise ExtensionError(
"ExtensionDefinition with id '{}' is not registered.".format(
extension_id))
del self.extensions[extension_id]
def add_extension(self, extension_definition):
e_id = extension_definition.extension_id
if e_id in self.extensions:
raise ExtensionError(
"ExtensionDefinition with id '{}' already exists.".format(
e_id))
self.extensions[e_id] = extension_definition
def __getitem__(self, extension_id):
"""Gets the extension object for the given extension."""
# Check to make sure this is a registered extension.
if not pystac_client.STAC_API_EXTENSIONS.is_registered_extension(
extension_id):
raise ExtensionError(
"'{}' is not an extension "
"registered with pystac-client".format(extension_id))
if not self.implements(extension_id):
raise ExtensionError(
"{} does not implement the {} extension. "
"Use the 'ext.enable' method to enable this extension "
"first.".format(self.stac_object, extension_id))
return pystac_client.STAC_API_EXTENSIONS.extend_object(
extension_id, self.stac_object)
def __init__(self, stac_object_class, extension_class):
if stac_object_class is Catalog:
if not issubclass(extension_class, CatalogExtension):
raise ExtensionError(
"Classes extending catalogs must inheret from CatalogExtension"
)
if stac_object_class is Collection:
if not issubclass(extension_class, CollectionExtension):
raise ExtensionError(
"Classes extending collections must inheret from CollectionExtension"
)
if stac_object_class is Item:
if not issubclass(extension_class, ItemExtension):
raise ExtensionError(
"Classes extending item must inheret from ItemExtension")
self.stac_object_class = stac_object_class
self.extension_class = extension_class
def __getitem__(self, extension_id):
"""Gets the extension object for the given extension.
Returns:
CatalogExtension or CollectionExtension or ItemExtension: The extension object
through which you can access the extension functionality for the extension represented
by the extension_id.
"""
# Check to make sure this is a registered extension.
if not pystac.STAC_EXTENSIONS.is_registered_extension(extension_id):
raise ExtensionError("'{}' is not an extension "
"registered with PySTAC".format(extension_id))
if not self.implements(extension_id):
raise ExtensionError("{} does not implement the {} extension. "
"Use the 'ext.enable' method to enable this extension "
"first.".format(self.stac_object, extension_id))
return pystac.STAC_EXTENSIONS.extend_object(extension_id, self.stac_object)
def extend_object(self, extension_id, stac_object):
"""Returns the extension object for the given STACObject and the given
extension_id
"""
ext_class = self.get_extension_class(extension_id, type(stac_object))
if ext_class is None:
raise ExtensionError(
"Extension '{}' does not extend objects of type {}".format(
extension_id, type(stac_object)))
return ext_class._from_object(stac_object)
def __init__(self, stac_object_class, extension_class):
if stac_object_class is API:
if not issubclass(extension_class, APIExtension):
raise ExtensionError(
"Classes extending API instances must inherit from APIExtension"
)
if stac_object_class is ItemSearch:
if not issubclass(extension_class, ItemSearchFragment):
raise ExtensionError(
"Classes extending ItemSearch instances must inherit from ItemSearchExtension"
)
if stac_object_class is ItemCollection:
if not issubclass(extension_class, ItemCollectionFragment):
raise ExtensionError(
"Classes extending ItemCollection instances must inherit from ItemCollectionExtension"
)
if stac_object_class is Catalog: # pragma: no cover
if not issubclass(extension_class, CatalogExtension):
raise ExtensionError(
"Classes extending catalogs must inherit from CatalogExtension"
)
if stac_object_class is Collection: # pragma: no cover
if not issubclass(extension_class, CollectionExtension):
raise ExtensionError(
"Classes extending collections must inherit from CollectionExtension"
)
if stac_object_class is Item: # pragma: no cover
if not issubclass(extension_class, ItemExtension):
raise ExtensionError(
"Classes extending item must inherit from ItemExtension")
self.stac_object_class = stac_object_class
self.extension_class = extension_class
def extend_object(self, stac_object, extension_id):
ext = self.extensions.get(extension_id)
if ext is None:
raise ExtensionError(
"No ExtensionDefinition registered with id '{}'. "
"Is the extension ID correct, or are you forgetting to call "
"'add_extension' for a custom extension?".format(extension_id))
stac_object_class = type(stac_object)
ext_classes = [
e.extension_class for e in ext.extended_objects
if issubclass(stac_object_class, e.stac_object_class)
]
ext_class = None
if len(ext_classes) == 0:
raise ExtensionError(
"Extension '{}' does not extend objects of type {}".format(
extension_id, ext_class))
elif len(ext_classes) == 1:
ext_class = ext_classes[0]
else:
# Need to check collection extensions before catalogs.
sort_key = {}
for c in ext_classes:
for i, base_cl in enumerate(
[ItemExtension, CollectionExtension, CatalogExtension]):
if issubclass(c, base_cl):
sort_key[c] = i
break
if c not in sort_key:
sort_key[c] = -1
ext_class = sorted(ext_classes, key=lambda c: sort_key[c])[0]
return ext_class._from_object(stac_object)
def enable_extension(self, extension_id, stac_object):
"""Enables the extension for the given object.
This will at least ensure the extension ID is in the object's "stac_extensions"
property. Some extensions may make additional modifications to the object.
"""
# Check to make sure this is a registered extension.
if not self.is_registered_extension(extension_id):
raise ExtensionError("'{}' is not an extension "
"registered with PySTAC".format(extension_id))
ext_class = self.get_extension_class(extension_id, type(stac_object))
if ext_class is None:
raise ExtensionError(
"Extension '{}' does not extend objects of type {}".format(
extension_id, type(stac_object)))
if stac_object.stac_extensions is None:
stac_object.stac_extensions = [extension_id]
elif extension_id not in stac_object.stac_extensions:
stac_object.stac_extensions.append(extension_id)
ext_class.enable_extension(stac_object)
def can_extend(self, extension_id, stac_object_class):
"""Returns True if the extension can extend the given object type.
Args:
extension_id (str): The extension ID to check.
stac_object_class: the class of the object to check. Will check against subclasses,
so will return the correct result even if the object is a subclass of Catalog,
Collection or Item.
Returns:
bool
"""
ext = self.extensions.get(extension_id)
# Check to make sure this is a registered extension.
if ext is None:
raise ExtensionError(
"'{}' is not a registered extension".format(extension_id))
return any([
e.extension_class for e in ext.extended_objects
if issubclass(stac_object_class, e.stac_object_class)
])
def get_extension_class(self, extension_id, stac_object_class):
"""Gets the extension class for a given stac object class if one exists, otherwise
returns None
"""
ext = self.extensions.get(extension_id)
if ext is None:
raise ExtensionError(
"No ExtensionDefinition registered with id '{}'. "
"Is the extension ID correct, or are you forgetting to call "
"'add_extension' for a custom extension?".format(extension_id))
ext_classes = [
e.extension_class for e in ext.extended_objects
if issubclass(stac_object_class, e.stac_object_class)
]
ext_class = None
if len(ext_classes) == 0:
return None
elif len(ext_classes) == 1:
ext_class = ext_classes[0]
else:
# Need to check collection extensions before catalogs.
sort_key = {}
for c in ext_classes:
for i, base_cl in enumerate(
[ItemExtension, CollectionExtension, CatalogExtension]):
if issubclass(c, base_cl):
sort_key[c] = i
break
if c not in sort_key:
sort_key[c] = -1
ext_class = sorted(ext_classes, key=lambda c: sort_key[c])[0]
return ext_class
