Skip to main content
`Ctrl`+`K`

# 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 f"<OrderEvent type={self.event_type} 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:

1. Make sure that your Extension class has a `name` attribute with `Literal(<name>)` as the type.

2. Import your Extension class in `pystac/extensions/ext.py`

3. Add the `name` to `EXTENSION_NAMES`

4. Add the mapping from name to class to `EXTENSION_NAME_MAPPING`

5. 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)