validate-stac-geoparquet

#!/usr/bin/env python

import json
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

import pyarrow
import pyarrow.parquet
import stac_geoparquet.arrow
from deepdiff import DeepDiff

root = Path(__file__).parents[1]
path = root / "spec-examples" / "v1.1.0" / "extended-item.json"
directory = tempfile.mkdtemp()
parquet_path = Path(directory) / "extended-item.parquet"


def clean_item(item: dict[str, Any]) -> None:
    if "type" not in item:
        item["type"] = "Feature"
    if (
        item["geometry"]["type"] == "MultiPolygon"
        and len(item["geometry"]["coordinates"]) == 1
    ):
        item["geometry"]["type"] = "Polygon"
        item["geometry"]["coordinates"] = item["geometry"]["coordinates"][0]


def clean_report(report: dict[str, Any]) -> dict[str, Any]:
    """We expect datetime values to be changed in the report."""
    if report.get("values_changed"):
        if report["values_changed"].get("root['properties']['datetime']") == {
            "new_value": "2020-12-14T18:02:31.437Z",
            "old_value": "2020-12-14T18:02:31.437000Z",
        }:
            del report["values_changed"]["root['properties']['datetime']"]
        if report["values_changed"].get("root['properties']['created']") == {
            "new_value": "2020-12-15T01:48:13.725+00:00",
            "old_value": "2020-12-15T01:48:13.725Z",
        }:
            del report["values_changed"]["root['properties']['created']"]
        if report["values_changed"].get("root['properties']['updated']") == {
            "new_value": "2020-12-15T01:48:13.725+00:00",
            "old_value": "2020-12-15T01:48:13.725Z",
        }:
            del report["values_changed"]["root['properties']['updated']"]
        if not report["values_changed"]:
            del report["values_changed"]
    return report


try:
    # Writing
    subprocess.check_call(
        [
            "cargo",
            "run",
            "-p",
            "rustac",
            "-F",
            "duckdb",
            "--",
            "translate",
            path,
            parquet_path,
        ]
    )
    table = pyarrow.parquet.read_table(parquet_path)
    after = next(stac_geoparquet.arrow.stac_table_to_items(table))
    clean_item(after)
    with open(path) as f:
        before = json.load(f)
    report = DeepDiff(before, after).to_dict()
    report = clean_report(report)
    if report:
        print(json.dumps(report, indent=2))
        sys.exit(1)
    else:
        parquet_path.unlink()

    # Reading
    table = stac_geoparquet.arrow.parse_stac_items_to_arrow([before])
    stac_geoparquet.arrow.to_parquet(table, parquet_path)
    item_collection = json.loads(
        subprocess.check_output(
            [
                "cargo",
                "run",
                "-p",
                "rustac",
                "-F",
                "duckdb",
                "--",
                "translate",
                parquet_path,
            ]
        )
    )
    assert len(item_collection["features"]) == 1
    clean_item(
        item_collection["features"][0]
    )  # stac-geoparquet writes as a multi-polygon
    report = DeepDiff(before, item_collection["features"][0]).to_dict()
    report = clean_report(report)
    if report:
        print(json.dumps(report, indent=2))
        sys.exit(1)

finally:
    shutil.rmtree(directory)