asset_vrt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple, Optional, Set\n",
    "from urllib.parse import urlparse\n",
    "import os\n",
    "import attr\n",
    "import warnings\n",
    "\n",
    "from rio_tiler import io\n",
    "from rio_tiler.io.stac import DEFAULT_VALID_TYPE\n",
    "from rio_tiler.types import AssetInfo\n",
    "from rio_tiler.errors import InvalidAssetName\n",
    "from rasterio._path import _vsi_path, _parse_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "STAC_ALTERNATE_KEY = os.environ.get(\"RIO_TILER_STAC_ALTERNATE_KEY\", None)\n",
    "\n",
    "VALID_TYPE = {\n",
    "    *DEFAULT_VALID_TYPE,\n",
    "    \"application/wmo-GRIB2\",\n",
    "}\n",
    "\n",
    "\n",
    "def _parse_asset(asset_string: str) -> Tuple[str, Optional[str]]:\n",
    "    if asset_string.startswith(\"vrt://\"):\n",
    "        parsed = urlparse(asset_string)\n",
    "\n",
    "        assert parsed.netloc, f\"Could not find asset name in {asset_string}\"\n",
    "        asset = parsed.netloc\n",
    "\n",
    "        return asset, parsed.query\n",
    "\n",
    "    return asset_string, None\n",
    "\n",
    "\n",
    "@attr.s\n",
    "class STACReader(io.STACReader):\n",
    "\n",
    "    include_asset_types: Set[str] = attr.ib(default=VALID_TYPE)\n",
    "\n",
    "    def _get_asset_info(self, asset: str) -> AssetInfo:\n",
    "        \"\"\"Validate asset names and return asset's info.\n",
    "\n",
    "        Args:\n",
    "            asset (str): STAC asset name.\n",
    "\n",
    "        Returns:\n",
    "            AssetInfo: STAC asset info.\n",
    "\n",
    "        \"\"\"\n",
    "        asset, asset_vrt_option = _parse_asset(asset)\n",
    "\n",
    "\n",
    "        if asset not in self.assets:\n",
    "            raise InvalidAssetName(\n",
    "                f\"'{asset}' is not valid, should be one of {self.assets}\"\n",
    "            )\n",
    "\n",
    "        asset_info = self.item.assets[asset]\n",
    "        extras = asset_info.extra_fields\n",
    "\n",
    "        info = AssetInfo(\n",
    "            url=asset_info.get_absolute_href() or asset_info.href,\n",
    "            metadata=extras,\n",
    "        )\n",
    "\n",
    "        if STAC_ALTERNATE_KEY and extras.get(\"alternate\"):\n",
    "            if alternate := extras[\"alternate\"].get(STAC_ALTERNATE_KEY):\n",
    "                info[\"url\"] = alternate[\"href\"]\n",
    "\n",
    "        if asset_info.media_type:\n",
    "            info[\"media_type\"] = asset_info.media_type\n",
    "\n",
    "        # https://github.com/stac-extensions/file\n",
    "        if head := extras.get(\"file:header_size\"):\n",
    "            info[\"env\"] = {\"GDAL_INGESTED_BYTES_AT_OPEN\": head}\n",
    "\n",
    "        # https://github.com/stac-extensions/raster\n",
    "        if bands := extras.get(\"raster:bands\"):\n",
    "            stats = [\n",
    "                (b[\"statistics\"][\"minimum\"], b[\"statistics\"][\"maximum\"])\n",
    "                for b in bands\n",
    "                if {\"minimum\", \"maximum\"}.issubset(b.get(\"statistics\", {}))\n",
    "            ]\n",
    "            # check that stats data are all double and make warning if not\n",
    "            if (\n",
    "                stats\n",
    "                and all(isinstance(v, (int, float)) for stat in stats for v in stat)\n",
    "                and len(stats) == len(bands)\n",
    "            ):\n",
    "                info[\"dataset_statistics\"] = stats\n",
    "            else:\n",
    "                warnings.warn(\n",
    "                    \"Some statistics data in STAC are invalid, they will be ignored.\"\n",
    "                )\n",
    "\n",
    "        if asset_vrt_option:\n",
    "            url = _vsi_path(_parse_path(info[\"url\"]))\n",
    "            info[\"url\"] = f\"vrt://{url}?{asset_vrt_option}\"\n",
    "\n",
    "        return info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STACReader(bounds=[-180.0, -90.0, 180.0, 90.0], crs=CRS.from_epsg(4326), transform=None, height=None, width=None, input='https://planetarycomputer.microsoft.com/api/stac/v1/collections/ecmwf-forecast/items/ecmwf-2024-09-10T12-waef-ep-360h-0.25', item=<Item id=ecmwf-2024-09-10T12-waef-ep-360h-0.25>, tms=<TileMatrixSet title='Google Maps Compatible for the World' id='WebMercatorQuad' crs='http://www.opengis.net/def/crs/EPSG/0/3857>, minzoom=0, maxzoom=24, geographic_crs=CRS.from_epsg(4326), include_assets=None, exclude_assets=None, exclude_asset_types=None, assets=['data'], default_assets=None, reader=<class 'rio_tiler.io.rasterio.Reader'>, reader_options={}, fetch_options={}, ctx=<class 'rasterio.env.Env'>, include_asset_types={'image/tiff; application=geotiff', 'image/tiff; application=geotiff; profile=cloud-optimized', 'application/x-hdf', 'image/vnd.stac.geotiff; cloud-optimized=true', 'image/tiff; profile=cloud-optimized; application=geotiff', 'image/jp2', 'application/x-hdf5', 'application/wmo-GRIB2', 'image/tiff', 'image/x.geotiff'})\n",
      "['data']\n",
      "{'url': 'https://ai4edataeuwest.blob.core.windows.net/ecmwf/20240910/12z/ifs/0p25/waef/20240910120000-360h-waef-ep.grib2', 'metadata': {}, 'media_type': 'application/wmo-GRIB2'}\n",
      "{'url': 'vrt:///vsicurl/https://ai4edataeuwest.blob.core.windows.net/ecmwf/20240910/12z/ifs/0p25/waef/20240910120000-360h-waef-ep.grib2?bands=1,2,3', 'metadata': {}, 'media_type': 'application/wmo-GRIB2'}\n"
     ]
    }
   ],
   "source": [
    "with STACReader(\"https://planetarycomputer.microsoft.com/api/stac/v1/collections/ecmwf-forecast/items/ecmwf-2024-09-10T12-waef-ep-360h-0.25\") as stac:\n",
    "    print(stac)\n",
    "    print(stac.assets)\n",
    "    print(stac._get_asset_info(\"data\"))\n",
    "    print(stac._get_asset_info(\"vrt://data?bands=1,2,3\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py39",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}