stac_pydantic benchmarks

test_stac_pydantic_benchmark.py

import copy
import json
import orjson

from dateutil import parser
from pystac import Asset, Item, Link
from stac_pydantic import Item as PydanticItem

self_link_dict = {
    "rel": "self",
    "type": "application/geo+json",
    "href": "https://my.stac/v0/collections/mycollection/items/my_item",
}

collection_link_dict = {
    "rel": "collection",
    "type": "application/json",
    "href": "https://my.stac/v0/collections/mycollection",
}

parent_link_dict = {
    "rel": "parent",
    "type": "application/json",
    "href": "https://my.stac/v0/collections/mycollection",
}


item_dict = {
    "type": "Feature",
    "stac_version": "1.0.0",
    "id": "my_item",
    "collection": "mycollection",
    "geometry": {
        "coordinates": [
            [
                [149.19940003830305, -36.428645762574625],
                [149.05104951177714, -39.259123256881225],
                [152.81379005469256, -39.32599372427865],
                [152.82080633518694, -36.489064495233244],
                [149.19940003830305, -36.428645762574625],
            ]
        ],
        "type": "Polygon",
    },
    "bbox": [
        149.05104951177714,
        -39.32599372427865,
        152.82080633518694,
        -36.428645762574625,
    ],
    "properties": {
        "datetime": "2022-11-21T19:09:07.587000Z",
        "created": "2023-08-03T00:31:23.346946Z",
        "updated": "2023-08-03T00:31:23.346946Z",
    },
    "stac_extensions": [],
    "links": [self_link_dict, collection_link_dict, parent_link_dict],
    "assets": {
        "data": {
            "href": "https://my.stac/v0/collections/mycollection/items/my_item/data.tif",
            "title": "Data",
            "media_type": "image/tiff; application=geotiff",
            "roles": ["data"],
        },
        "thumbnail": {
            "href": "https://my.stac/v0/collections/mycollection/items/my_item/thumbnail.png",
            "media_type": "image/png",
            "roles": ["thumbnail"],
        },
        "metadata": {
            "href": "https://my.stac/v0/collections/mycollection/items/my_item/metadata.json",
            "media_type": "application/json",
            "roles": ["metadata"],
        },
    },
}

item_json: bytes = orjson.dumps(item_dict)

item_dict_2 = copy.deepcopy(item_dict)
pydantic_item = PydanticItem.model_validate(item_dict_2)

pystac_datetime = parser.parse(item_dict["properties"]["datetime"])


def create_pystac_item() -> Item:
    item = Item(
        id=item_dict["id"],
        geometry=item_dict["geometry"],
        bbox=item_dict["bbox"],
        datetime=pystac_datetime,
        properties={
            "created": item_dict["properties"]["created"],
            "updated": item_dict["properties"]["updated"],
        },
        collection=item_dict["collection"],
    )

    data_asset = Asset(**item_dict["assets"]["data"])
    thumbnail_asset = Asset(**item_dict["assets"]["thumbnail"])
    metadata_asset = Asset(**item_dict["assets"]["metadata"])

    item.add_asset("data", data_asset)
    item.add_asset("thumbnail", thumbnail_asset)
    item.add_asset("metadata", metadata_asset)

    item.set_self_href(self_link_dict["href"])
    collection_link = Link(
        collection_link_dict["rel"],
        target=collection_link_dict["href"],
        media_type=collection_link_dict["type"],
    )
    parent_link = Link(
        parent_link_dict["rel"],
        target=parent_link_dict["href"],
        media_type=parent_link_dict["type"],
    )

    item.add_link(collection_link)
    item.add_link(parent_link)

    return item


def create_pystac_item_validate() -> Item:
    item = create_pystac_item()
    item.validate()
    return item


pystac_item = create_pystac_item_validate()


def dump_dict_json() -> str:
    return json.dumps(item_dict)


def load_dict_json() -> dict:
    return json.loads(item_json)


def dump_dict_orjson() -> str:
    return orjson.dumps(item_dict).decode()


def load_dict_orjson() -> dict:
    return orjson.loads(item_json)


def load_pydantic_model_validate() -> PydanticItem:
    return PydanticItem.model_validate(item_dict_2)


def load_pydantic_model_validate_json() -> PydanticItem:
    return PydanticItem.model_validate_json(item_json)


def load_pydantic_orjson_model_validate() -> PydanticItem:
    return PydanticItem.model_validate(orjson.loads(item_json))


def load_pydantic_model_construct() -> PydanticItem:
    return PydanticItem.model_construct(item_dict_2)


def load_pydantic_orjson_model_construct() -> PydanticItem:
    return PydanticItem.model_construct(orjson.loads(item_json))


def pydantic_model_dump_json() -> str:
    return pydantic_item.model_dump_json()


def dump_pystac_json() -> str:
    return json.dumps(pystac_item.to_dict())


def dump_pystac_orjson() -> str:
    return orjson.dumps(pystac_item.to_dict()).decode()


class TestDump:
    def test_dict_json(self, benchmark):
        r = benchmark(dump_dict_json)

    def test_dict_orjson(self, benchmark):
        r = benchmark(dump_dict_orjson)

    def test_pydantic_model_dump_json(self, benchmark):
        r = benchmark(pydantic_model_dump_json)

    def test_pystac_json(self, benchmark):
        r = benchmark(dump_pystac_json)

    def test_pystac_orjson(self, benchmark):
        r = benchmark(dump_pystac_orjson)


class TestLoad:
    def test_dict_json(self, benchmark):
        r = benchmark(load_dict_json)

    def test_dict_orjson(self, benchmark):
        r = benchmark(load_dict_orjson)

    def test_pydantic_model_validate(self, benchmark):
        r = benchmark(load_pydantic_model_validate)

    def test_pydantic_model_validate_json(self, benchmark):
        r = benchmark(load_pydantic_model_validate_json)

    def test_pydantic_orjson_model_validate(self, benchmark):
        r = benchmark(load_pydantic_orjson_model_validate)

    def test_pydantic_model_construct(self, benchmark):
        r = benchmark(load_pydantic_model_construct)

    def test_pydantic_orjson_model_construct(self, benchmark):
        r = benchmark(load_pydantic_orjson_model_construct)

    def test_pystac_item(self, benchmark):
        r = benchmark(create_pystac_item)

    def test_pystac_item_validate(self, benchmark):
        r = benchmark(create_pystac_item_validate)