Adding New and Custom Extensions#

This tutorial will cover using the PropertiesExtension and ExtensionManagementMixin classes in pystac.extensions.base to implement a new extension in PySTAC, and how to make that class accessible via the pystac.Item.ext interface.

For this exercise, we will implement an imaginary Order Request Extension that allows us to track an internal order ID associated with a given satellite image, as well as the history of that imagery order. This use-case is specific enough that it would probably not be a good candidate for an actual STAC Extension, but it gives us an opportunity to highlight some of the key aspects and patterns used in implementing STAC Extensions in PySTAC.

First, we import the PySTAC modules and classes that we will be using throughout the tutorial.

[1]:

from datetime import datetime, timedelta
from pprint import pprint
from typing import Any, Dict, List, Literal, Optional, Union
from uuid import uuid4

import pystac
from pystac.extensions.base import ExtensionManagementMixin, PropertiesExtension
from pystac.utils import (
    StringEnum,
    datetime_to_str,
    get_required,
    map_opt,
    str_to_datetime,
)

Define the Extension#

Our extension will extend STAC Items by adding the following properties:

order:id: A unique string ID associated with the internal order for this image. This field will be required.
order:history: A chronological list of events associated with this order. Each of these “events” will have a timestamp and an event type, which will be one of the following: submitted, started_processing, delivered, cancelled. This field will be optional.

Create Extension Classes#

Let’s start by creating a class to represent the order history events.

[2]:

class OrderEventType(StringEnum):
    SUBMITTED = "submitted"
    STARTED_PROCESSING = "started_processing"
    DELIVERED = "delivered"
    CANCELLED = "cancelled"


class OrderEvent:
    properties: Dict[str, Any]

    def __init__(self, properties: Dict[str, Any]) -> None:
        self.properties = properties

    @property
    def event_type(self) -> OrderEventType:
        return get_required(self.properties.get("type"), self, "event_type")

    @event_type.setter
    def event_type(self, v: OrderEventType) -> None:
        self.properties["type"] = str(v)

    @property
    def timestamp(self) -> datetime:
        return str_to_datetime(
            get_required(self.properties.get("timestamp"), self, "timestamp")
        )

    @timestamp.setter
    def timestamp(self, v: datetime) -> None:
        self.properties["timestamp"] = datetime_to_str(v)

    def __repr__(self) -> str:
        return "<OrderEvent " f"type={self.event_type} " f"timestamp={self.timestamp}>"

    def apply(
        self,
        event_type: OrderEventType,
        timestamp: datetime,
    ) -> None:
        self.event_type = event_type
        self.timestamp = timestamp

    @classmethod
    def create(
        cls,
        event_type: OrderEventType,
        timestamp: datetime,
    ) -> "OrderEvent":
        oe = cls({})
        oe.apply(event_type=event_type, timestamp=timestamp)
        return oe

    def to_dict(self) -> Dict[str, Any]:
        return self.properties

A few important notes about how we constructed this:

We used PySTAC’s StringEnum class, which inherits from the Python Enum class, to capture the allowed event type values. This class has built-in methods that will convert these instances to strings when serializing STAC Items to JSON.
We use property getters and setters to manipulate a properties dictionary in our OrderEvent class. We will see later how this pattern allows us to mutate Item property dictionaries in-place so that updates to the OrderEvent object are synced to the Item they extend.
The timestamp property is converted to a string before it is saved in the properties dictionary. This ensures that dictionary is always JSON-serializable but allows us to work with the values as a Python datetime instance when using the property getter.
We use event_type as our property name so that we do not shadow the built-in type function in the apply method. However, this values is stored under the desired "type" key in the underlying properties dictionary.

Next, we will create a new class inheriting from PropertiesExtension and ExtensionManagementMixin. Since this class only extends pystac.Item instance, we do not need to make it generic. If you were creating an extension that applied to multiple object types (e.g. pystac.Item and pystac.Asset) then you would need to inherit from typing.Generic as well and create concrete extension classed for each of these object types (see the EOExtension, ItemEOExtension, and AssetEOExtension classes for an example of this implementation).

[3]:

SCHEMA_URI: str = "https://example.com/image-order/v1.0.0/schema.json"
PREFIX: str = "order:"
ID_PROP: str = PREFIX + "id"
HISTORY_PROP: str = PREFIX + "history"


class OrderExtension(
    PropertiesExtension, ExtensionManagementMixin[Union[pystac.Item, pystac.Collection]]
):
    name: Literal["order"] = "order"

    def __init__(self, item: pystac.Item):
        self.item = item
        self.properties = item.properties

    def apply(
        self, order_id: str = None, history: Optional[List[OrderEvent]] = None
    ) -> None:
        self.order_id = order_id
        self.history = history

    @property
    def order_id(self) -> str:
        return get_required(self._get_property(ID_PROP, str), self, ID_PROP)

    @order_id.setter
    def order_id(self, v: str) -> None:
        self._set_property(ID_PROP, v, pop_if_none=False)

    @property
    def history(self) -> Optional[List[OrderEvent]]:
        return map_opt(
            lambda history: [OrderEvent(d) for d in history],
            self._get_property(HISTORY_PROP, List[OrderEvent]),
        )

    @history.setter
    def history(self, v: Optional[List[OrderEvent]]) -> None:
        self._set_property(
            HISTORY_PROP,
            map_opt(lambda history: [event.to_dict() for event in history], v),
            pop_if_none=True,
        )

    @classmethod
    def get_schema_uri(cls) -> str:
        return SCHEMA_URI

    @classmethod
    def ext(cls, obj: pystac.Item, add_if_missing: bool = False) -> "OrderExtension":
        if isinstance(obj, pystac.Item):
            cls.ensure_has_extension(obj, add_if_missing)
            return OrderExtension(obj)
        else:
            raise pystac.ExtensionTypeError(
                f"OrderExtension does not apply to type '{type(obj).__name__}'"
            )

As with the OrderEvent class, we use property getters and setters for our extension fields (the PropertiesExtension class has a properties attribute where these are stored). Rather than setting these values directly in the dictionary, we use the _get_property and _set_property methods that are built into the PropertiesExtension class). We also add an ext method that will be used to extend pystac.Item instances, and a get_schema_uri method that is required for all PropertiesExtension classes.

Use the Extension#

Let’s try using our new classes to extend an Item and access the extension properties. We’ll start by loading the core Item example from the STAC spec examples here and printing the existing properties.

[4]:

item = pystac.read_file(
    "https://raw.githubusercontent.com/radiantearth/stac-spec/master/examples/core-item.json"
)
item.properties

[4]:

{'title': 'Core Item',
 'description': 'A sample STAC Item that includes examples of all common metadata',
 'datetime': None,
 'start_datetime': '2020-12-11T22:38:32.125Z',
 'end_datetime': '2020-12-11T22:38:32.327Z',
 'created': '2020-12-12T01:48:13.725Z',
 'updated': '2020-12-12T01:48:13.725Z',
 'platform': 'cool_sat1',
 'instruments': ['cool_sensor_v1'],
 'constellation': 'ion',
 'mission': 'collection 5624',
 'gsd': 0.512}

Next, let’s verify that this Item does not implement our new Order Extension yet and that it does not already contain any of our Order Extension properties.

[5]:

print(f"Implements Extension: {OrderExtension.has_extension(item)}")
print(f"Order ID: {item.properties.get(ID_PROP)}")
print("History:")
for event in item.properties.get(HISTORY_PROP, []):
    pprint(event)

Implements Extension: False
Order ID: None
History:

As expected, this Item does not implement the extension (i.e. the schema URI is not in the Item’s stac_extensions list). Let’s add it, create an instance of OrderExtension that extends the Item, and add some values for our extension fields.

[6]:

order_ext = OrderExtension.ext(item, add_if_missing=True)

# Create a unique string ID for the order ID
order_ext.order_id = str(uuid4())

# Create some fake order history and set it using the extension
event_1 = OrderEvent.create(
    event_type=OrderEventType.SUBMITTED, timestamp=datetime.now() - timedelta(days=1)
)
event_2 = OrderEvent.create(
    event_type=OrderEventType.STARTED_PROCESSING,
    timestamp=datetime.now() - timedelta(hours=12),
)
event_3 = OrderEvent.create(
    event_type=OrderEventType.DELIVERED, timestamp=datetime.now() - timedelta(hours=1)
)
order_ext.history = [event_1, event_2, event_3]

Now let’s check to see if these values were written to our Item properties.

[7]:

print(f"Implements Extension: {OrderExtension.has_extension(item)}")
print(f"Order ID: {item.properties.get(ID_PROP)}")
print("History:")
for event in item.properties.get(HISTORY_PROP, []):
    pprint(event)

Implements Extension: True
Order ID: 7a206229-78f0-46cb-afc2-acf45e14afab
History:
{'timestamp': '2023-10-11T11:21:50.989315Z', 'type': 'submitted'}
{'timestamp': '2023-10-11T23:21:50.989372Z', 'type': 'started_processing'}
{'timestamp': '2023-10-12T10:21:50.989403Z', 'type': 'delivered'}

(Optional) Add access via `Item.ext`#

This applies if you are planning on opening a Pull Request to add this implementation of the extension class to the pystac library

Now that you have a complete extension class, you can add access to it via the pystac.Item.ext interface by following these steps:

Make sure that your Extension class has a name attribute with Literal(<name>) as the type.
Import your Extension class in pystac/extensions/ext.py
Add the name to EXTENSION_NAMES
Add the mapping from name to class to EXTENSION_NAME_MAPPING
Add a getter method to the Ext class for any object type that this extension works with.

Here is an example of the diff:

diff --git a/pystac/extensions/ext.py b/pystac/extensions/ext.py
index 93a30fe..2dbe5ca 100644
--- a/pystac/extensions/ext.py
+++ b/pystac/extensions/ext.py
@@ -9,6 +9,7 @@ from pystac.extensions.file import FileExtension
 from pystac.extensions.grid import GridExtension
 from pystac.extensions.item_assets import AssetDefinition, ItemAssetsExtension
 from pystac.extensions.mgrs import MgrsExtension
+from pystac.extensions.order import OrderExtension
 from pystac.extensions.pointcloud import PointcloudExtension
 from pystac.extensions.projection import ProjectionExtension
 from pystac.extensions.raster import RasterExtension
@@ -32,6 +33,7 @@ EXTENSION_NAMES = Literal[
     "grid",
     "item_assets",
     "mgrs",
+    "order",
     "pc",
     "proj",
     "raster",
@@ -54,6 +56,7 @@ EXTENSION_NAME_MAPPING: Dict[EXTENSION_NAMES, Any] = {
     GridExtension.name: GridExtension,
     ItemAssetsExtension.name: ItemAssetsExtension,
     MgrsExtension.name: MgrsExtension,
+    OrderExtension.name: OrderExtension,
     PointcloudExtension.name: PointcloudExtension,
     ProjectionExtension.name: ProjectionExtension,
     RasterExtension.name: RasterExtension,
@@ -150,6 +153,10 @@ class ItemExt:
     def mgrs(self) -> MgrsExtension:
         return MgrsExtension.ext(self.stac_object)

+    @property
+    def order(self) -> OrderExtension:
+        return OrderExtension.ext(self.stac_object)
+
     @property
     def pc(self) -> PointcloudExtension[pystac.Item]:
         return PointcloudExtension.ext(self.stac_object)

Adding New and Custom Extensions#

Define the Extension#

Create Extension Classes#

Use the Extension#

(Optional) Add access via Item.ext#

(Optional) Add access via `Item.ext`#