tilefetch 0.1.0__py3-none-any.whl

tilefetch/pmtiles.py

@@ -0,0 +1,657 @@
+"""PMTiles version 3 writer and reader (pure Python, stdlib only).
+
+Implements the layout of the PMTiles v3 specification
+(https://github.com/protomaps/PMTiles/blob/main/spec/v3/spec.md):
+
+* a fixed 127-byte little-endian header;
+* directories of ``(tile_id, offset, length, run_length)`` entries, serialised
+  as unsigned LEB128 varints in four columns (delta-coded tile IDs, run
+  lengths, lengths, and offsets where ``0`` means "directly after the previous
+  entry") and gzip-compressed;
+* tile IDs that walk zoom levels in order and the Hilbert curve within a level;
+* a root directory that fits, together with the header, in the first 16 KiB,
+  with leaf directories when needed;
+* de-duplication of identical tiles and run-length encoding of repeats.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import gzip
+import hashlib
+import io
+import json
+import os
+import sqlite3
+import struct
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import IO, Any, Iterator, NamedTuple
+
+from .tiles import Tile
+
+__all__ = [
+    "HEADER_SIZE",
+    "Compression",
+    "TileType",
+    "Entry",
+    "Header",
+    "encode_varint",
+    "decode_varint",
+    "zxy_to_tileid",
+    "tileid_to_zxy",
+    "serialize_directory",
+    "deserialize_directory",
+    "PMTilesWriter",
+    "PMTilesReader",
+    "write_pmtiles",
+]
+
+HEADER_SIZE = 127
+_ROOT_LIMIT = 16384
+_MIN_LEAF_SIZE = 4096
+_MAGIC = b"PMTiles"
+_HEADER_FMT = "<7sB11Q6B4iB2i"
+assert struct.calcsize(_HEADER_FMT) == HEADER_SIZE
+
+
+class Compression:
+    UNKNOWN = 0
+    NONE = 1
+    GZIP = 2
+    BROTLI = 3
+    ZSTD = 4
+
+
+class TileType:
+    UNKNOWN = 0
+    MVT = 1
+    PNG = 2
+    JPEG = 3
+    WEBP = 4
+    AVIF = 5
+
+
+def tile_type_of(data: bytes) -> int:
+    """Guess the PMTiles tile type from a tile's magic bytes."""
+    if data.startswith(b"\x89PNG\r\n\x1a\n"):
+        return TileType.PNG
+    if data.startswith(b"\xff\xd8\xff"):
+        return TileType.JPEG
+    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
+        return TileType.WEBP
+    if data[4:12] in (b"ftypavif", b"ftypavis"):
+        return TileType.AVIF
+    return TileType.UNKNOWN
+
+
+# ---------------------------------------------------------------------------
+# Varints and tile IDs
+# ---------------------------------------------------------------------------
+
+
+def encode_varint(value: int, out: bytearray | None = None) -> bytearray:
+    """Append ``value`` as an unsigned LEB128 varint to ``out`` and return it."""
+    if value < 0:
+        raise ValueError("varints are unsigned")
+    buf = out if out is not None else bytearray()
+    while True:
+        byte = value & 0x7F
+        value >>= 7
+        if value:
+            buf.append(byte | 0x80)
+        else:
+            buf.append(byte)
+            return buf
+
+
+def decode_varint(data: bytes, pos: int) -> tuple[int, int]:
+    """Decode one unsigned LEB128 varint; returns ``(value, new_pos)``."""
+    result = 0
+    shift = 0
+    while True:
+        if pos >= len(data):
+            raise ValueError("truncated varint")
+        b = data[pos]
+        pos += 1
+        result |= (b & 0x7F) << shift
+        if not b & 0x80:
+            return result, pos
+        shift += 7
+        if shift > 63:
+            raise ValueError("varint is too long")
+
+
+def _rotate(n: int, x: int, y: int, rx: int, ry: int) -> tuple[int, int]:
+    if ry == 0:
+        if rx == 1:
+            x = n - 1 - x
+            y = n - 1 - y
+        return y, x
+    return x, y
+
+
+def zxy_to_tileid(z: int, x: int, y: int) -> int:
+    """PMTiles tile ID: tiles of all lower zooms, plus the Hilbert index at ``z``."""
+    if z < 0 or z > 31:
+        raise ValueError("zoom must be between 0 and 31")
+    n = 1 << z
+    if not (0 <= x < n and 0 <= y < n):
+        raise ValueError(f"tile {z}/{x}/{y} is outside the zoom-{z} grid")
+    acc = ((1 << (2 * z)) - 1) // 3
+    d = 0
+    s = n >> 1
+    while s > 0:
+        rx = 1 if x & s else 0
+        ry = 1 if y & s else 0
+        d += s * s * ((3 * rx) ^ ry)
+        x, y = _rotate(n, x, y, rx, ry)
+        s >>= 1
+    return acc + d
+
+
+def tileid_to_zxy(tile_id: int) -> Tile:
+    """Inverse of :func:`zxy_to_tileid`."""
+    if tile_id < 0:
+        raise ValueError("tile id must be non-negative")
+    acc = 0
+    for z in range(32):
+        count = 1 << (2 * z)
+        if tile_id < acc + count:
+            d = tile_id - acc
+            n = 1 << z
+            x = y = 0
+            s = 1
+            while s < n:
+                rx = 1 & (d // 2)
+                ry = 1 & (d ^ rx)
+                x, y = _rotate(s, x, y, rx, ry)
+                x += s * rx
+                y += s * ry
+                d //= 4
+                s <<= 1
+            return Tile(z, x, y)
+        acc += count
+    raise ValueError("tile id is too large")
+
+
+# ---------------------------------------------------------------------------
+# Directories
+# ---------------------------------------------------------------------------
+
+
+class Entry(NamedTuple):
+    tile_id: int
+    offset: int
+    length: int
+    run_length: int
+
+
+def serialize_directory(entries: list[Entry]) -> bytes:
+    """Serialise directory entries (uncompressed) per the v3 spec."""
+    buf = bytearray()
+    encode_varint(len(entries), buf)
+    last = 0
+    for e in entries:
+        encode_varint(e.tile_id - last, buf)
+        last = e.tile_id
+    for e in entries:
+        encode_varint(e.run_length, buf)
+    for e in entries:
+        encode_varint(e.length, buf)
+    for i, e in enumerate(entries):
+        if i > 0 and e.offset == entries[i - 1].offset + entries[i - 1].length:
+            encode_varint(0, buf)
+        else:
+            encode_varint(e.offset + 1, buf)
+    return bytes(buf)
+
+
+def deserialize_directory(data: bytes) -> list[Entry]:
+    """Inverse of :func:`serialize_directory`."""
+    n, pos = decode_varint(data, 0)
+    ids = []
+    last = 0
+    for _ in range(n):
+        d, pos = decode_varint(data, pos)
+        last += d
+        ids.append(last)
+    runs = []
+    for _ in range(n):
+        v, pos = decode_varint(data, pos)
+        runs.append(v)
+    lengths = []
+    for _ in range(n):
+        v, pos = decode_varint(data, pos)
+        lengths.append(v)
+    entries: list[Entry] = []
+    for i in range(n):
+        v, pos = decode_varint(data, pos)
+        if v == 0 and i > 0:
+            offset = entries[i - 1].offset + entries[i - 1].length
+        else:
+            offset = v - 1
+        entries.append(Entry(ids[i], offset, lengths[i], runs[i]))
+    return entries
+
+
+def _compress(data: bytes, compression: int) -> bytes:
+    if compression == Compression.GZIP:
+        return gzip.compress(data, compresslevel=9, mtime=0)
+    if compression == Compression.NONE:
+        return data
+    raise ValueError(f"unsupported internal compression {compression}")
+
+
+def _decompress(data: bytes, compression: int) -> bytes:
+    if compression == Compression.GZIP:
+        return gzip.decompress(data)
+    if compression in (Compression.NONE, Compression.UNKNOWN):
+        return data
+    raise ValueError(f"unsupported internal compression {compression}")
+
+
+def _build_directories(entries: list[Entry], compression: int) -> tuple[bytes, bytes]:
+    """Return ``(root, leaves)``; the root fits in the first 16 KiB with the header."""
+    root = _compress(serialize_directory(entries), compression)
+    if len(root) <= _ROOT_LIMIT - HEADER_SIZE:
+        return root, b""
+    leaf_size = max(_MIN_LEAF_SIZE, len(entries) // 3500)
+    while True:
+        leaves = bytearray()
+        root_entries = []
+        for i in range(0, len(entries), leaf_size):
+            chunk = entries[i : i + leaf_size]
+            blob = _compress(serialize_directory(chunk), compression)
+            root_entries.append(Entry(chunk[0].tile_id, len(leaves), len(blob), 0))
+            leaves.extend(blob)
+        root = _compress(serialize_directory(root_entries), compression)
+        if len(root) <= _ROOT_LIMIT - HEADER_SIZE:
+            return root, bytes(leaves)
+        leaf_size *= 2
+
+
+# ---------------------------------------------------------------------------
+# Header
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Header:
+    root_offset: int = HEADER_SIZE
+    root_length: int = 0
+    metadata_offset: int = 0
+    metadata_length: int = 0
+    leaf_offset: int = 0
+    leaf_length: int = 0
+    data_offset: int = 0
+    data_length: int = 0
+    addressed_tiles: int = 0
+    tile_entries: int = 0
+    tile_contents: int = 0
+    clustered: bool = True
+    internal_compression: int = Compression.GZIP
+    tile_compression: int = Compression.NONE
+    tile_type: int = TileType.UNKNOWN
+    min_zoom: int = 0
+    max_zoom: int = 0
+    min_lon: float = -180.0
+    min_lat: float = -85.0
+    max_lon: float = 180.0
+    max_lat: float = 85.0
+    center_zoom: int = 0
+    center_lon: float = 0.0
+    center_lat: float = 0.0
+
+    def pack(self) -> bytes:
+        def e7(v: float) -> int:
+            return int(round(v * 10_000_000))
+
+        return struct.pack(
+            _HEADER_FMT,
+            _MAGIC,
+            3,
+            self.root_offset,
+            self.root_length,
+            self.metadata_offset,
+            self.metadata_length,
+            self.leaf_offset,
+            self.leaf_length,
+            self.data_offset,
+            self.data_length,
+            self.addressed_tiles,
+            self.tile_entries,
+            self.tile_contents,
+            1 if self.clustered else 0,
+            self.internal_compression,
+            self.tile_compression,
+            self.tile_type,
+            self.min_zoom,
+            self.max_zoom,
+            e7(self.min_lon),
+            e7(self.min_lat),
+            e7(self.max_lon),
+            e7(self.max_lat),
+            self.center_zoom,
+            e7(self.center_lon),
+            e7(self.center_lat),
+        )
+
+    @classmethod
+    def unpack(cls, data: bytes) -> Header:
+        if len(data) < HEADER_SIZE:
+            raise ValueError("file is too short to be a PMTiles archive")
+        v = struct.unpack(_HEADER_FMT, data[:HEADER_SIZE])
+        if v[0] != _MAGIC:
+            raise ValueError("not a PMTiles archive (bad magic)")
+        if v[1] != 3:
+            raise ValueError(f"unsupported PMTiles version {v[1]}")
+        values: list[Any] = [*v[2:13], bool(v[13]), *v[14:19]]
+        values += [
+            v[19] / 1e7,
+            v[20] / 1e7,
+            v[21] / 1e7,
+            v[22] / 1e7,
+            v[23],
+            v[24] / 1e7,
+            v[25] / 1e7,
+        ]
+        names = [f.name for f in dataclasses.fields(cls)]
+        return cls(**dict(zip(names, values)))
+
+
+# ---------------------------------------------------------------------------
+# Writer
+# ---------------------------------------------------------------------------
+
+
+class PMTilesWriter:
+    """Collect tiles, then write a PMTiles v3 archive on :meth:`close`.
+
+    Tiles are staged in an SQLite file next to the output (``<path>.part``), so
+    an interrupted download can resume: reopening the writer keeps the staged
+    tiles. The archive is assembled in tile-ID order (``clustered``) when the
+    writer is closed, and the staging file is then removed.
+    """
+
+    def __init__(
+        self, path: str | os.PathLike[str], metadata: dict[str, Any] | None = None
+    ) -> None:
+        self.path = os.fspath(path)
+        self.staging = self.path + ".part"
+        self.metadata: dict[str, Any] = dict(metadata or {})
+        self._db = sqlite3.connect(self.staging)
+        self._db.execute(
+            "CREATE TABLE IF NOT EXISTS tiles (tile_id INTEGER PRIMARY KEY, z INTEGER, data BLOB)"
+        )
+        self._db.execute("CREATE TABLE IF NOT EXISTS meta (k TEXT PRIMARY KEY, v TEXT)")
+        for k, v in self._db.execute("SELECT k, v FROM meta"):
+            self.metadata.setdefault(k, json.loads(v))
+        self._db.commit()
+        self._pending = 0
+        self.closed = False
+
+    def __enter__(self) -> PMTilesWriter:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        if exc[0] is None:
+            self.close()
+        else:
+            self.abort()
+
+    def has(self, tile: Tile) -> bool:
+        row = self._db.execute(
+            "SELECT 1 FROM tiles WHERE tile_id = ?", (zxy_to_tileid(*tile),)
+        ).fetchone()
+        return row is not None
+
+    def put(self, tile: Tile, data: bytes) -> None:
+        self._db.execute(
+            "INSERT OR REPLACE INTO tiles (tile_id, z, data) VALUES (?, ?, ?)",
+            (zxy_to_tileid(*tile), tile[0], sqlite3.Binary(data)),
+        )
+        self._pending += 1
+        if self._pending >= 256:
+            self._db.commit()
+            self._pending = 0
+
+    def set_metadata(self, metadata: dict[str, Any]) -> None:
+        self.metadata.update(metadata)
+        for k, v in metadata.items():
+            self._db.execute("INSERT OR REPLACE INTO meta (k, v) VALUES (?, ?)", (k, json.dumps(v)))
+
+    def __len__(self) -> int:
+        return int(self._db.execute("SELECT COUNT(*) FROM tiles").fetchone()[0])
+
+    def abort(self) -> None:
+        """Close without writing the archive; the staging file is kept for resume."""
+        if not self.closed:
+            self._db.commit()
+            self._db.close()
+            self.closed = True
+
+    def close(self) -> None:
+        """Write the archive and delete the staging file."""
+        if self.closed:
+            return
+        self._db.commit()
+        rows = self._db.execute("SELECT tile_id, data FROM tiles ORDER BY tile_id")
+        write_pmtiles(self.path, ((int(t), bytes(d)) for t, d in rows), self.metadata)
+        self._db.close()
+        self.closed = True
+        os.remove(self.staging)
+
+
+def write_pmtiles(
+    path: str | os.PathLike[str] | IO[bytes],
+    tiles: Any,
+    metadata: dict[str, Any] | None = None,
+    *,
+    internal_compression: int = Compression.GZIP,
+) -> Header:
+    """Write a PMTiles v3 archive.
+
+    ``tiles`` is an iterable of ``(tile_id, bytes)`` sorted by tile ID, or of
+    ``(Tile, bytes)`` in any order. ``metadata`` is stored as the JSON metadata
+    block; the keys ``bounds`` (``[w, s, e, n]``), ``center``
+    (``[lon, lat, zoom]``), ``minzoom`` and ``maxzoom`` also fill the header
+    when present.
+    """
+    meta = dict(metadata or {})
+    items = list(tiles)
+    if items and not isinstance(items[0][0], int):
+        items = sorted(((zxy_to_tileid(*t), d) for t, d in items), key=lambda it: it[0])
+    data_buf = io.BytesIO()
+    offsets: dict[bytes, tuple[int, int]] = {}
+    entries: list[Entry] = []
+    tile_type = TileType.UNKNOWN
+    last_id = -1
+    min_z, max_z = 99, -1
+    for tile_id, data in items:
+        if tile_id <= last_id:
+            raise ValueError("tile IDs must be unique and sorted")
+        last_id = tile_id
+        if tile_type == TileType.UNKNOWN:
+            tile_type = tile_type_of(data)
+        digest = hashlib.sha256(data).digest()
+        if digest in offsets:
+            off, length = offsets[digest]
+        else:
+            off, length = data_buf.tell(), len(data)
+            data_buf.write(data)
+            offsets[digest] = (off, length)
+        prev = entries[-1] if entries else None
+        if prev is not None and prev.offset == off and prev.tile_id + prev.run_length == tile_id:
+            entries[-1] = prev._replace(run_length=prev.run_length + 1)
+        else:
+            entries.append(Entry(tile_id, off, length, 1))
+        z = tileid_to_zxy(tile_id).z
+        min_z, max_z = min(min_z, z), max(max_z, z)
+    if not entries:
+        min_z = max_z = 0
+
+    root, leaves = _build_directories(entries, internal_compression)
+    meta_bytes = _compress(
+        json.dumps(meta, separators=(",", ":")).encode("utf-8"), internal_compression
+    )
+    tile_data = data_buf.getvalue()
+
+    h = Header()
+    h.root_offset = HEADER_SIZE
+    h.root_length = len(root)
+    h.metadata_offset = h.root_offset + h.root_length
+    h.metadata_length = len(meta_bytes)
+    h.leaf_offset = h.metadata_offset + h.metadata_length
+    h.leaf_length = len(leaves)
+    h.data_offset = h.leaf_offset + h.leaf_length
+    h.data_length = len(tile_data)
+    h.addressed_tiles = len(items)
+    h.tile_entries = len(entries)
+    h.tile_contents = len(offsets)
+    h.clustered = True
+    h.internal_compression = internal_compression
+    h.tile_compression = Compression.NONE if tile_type != TileType.MVT else Compression.UNKNOWN
+    h.tile_type = tile_type
+    h.min_zoom = int(meta.get("minzoom", min_z))
+    h.max_zoom = int(meta.get("maxzoom", max_z))
+    if "bounds" in meta:
+        h.min_lon, h.min_lat, h.max_lon, h.max_lat = (float(v) for v in meta["bounds"])
+    if "center" in meta:
+        h.center_lon, h.center_lat = float(meta["center"][0]), float(meta["center"][1])
+        h.center_zoom = int(meta["center"][2])
+    else:
+        h.center_lon = (h.min_lon + h.max_lon) / 2.0
+        h.center_lat = (h.min_lat + h.max_lat) / 2.0
+        h.center_zoom = h.min_zoom
+
+    blob = h.pack() + root + meta_bytes + leaves + tile_data
+    if hasattr(path, "write"):
+        path.write(blob)
+    else:
+        tmp = os.fspath(path) + ".tmp"
+        with open(tmp, "wb") as fh:
+            fh.write(blob)
+        os.replace(tmp, os.fspath(path))
+    return h
+
+
+# ---------------------------------------------------------------------------
+# Reader
+# ---------------------------------------------------------------------------
+
+
+def _find(entries: list[Entry], tile_id: int) -> Entry | None:
+    lo, hi = 0, len(entries) - 1
+    while lo <= hi:
+        mid = (lo + hi) // 2
+        c = entries[mid].tile_id
+        if c < tile_id:
+            lo = mid + 1
+        elif c > tile_id:
+            hi = mid - 1
+        else:
+            return entries[mid]
+    if hi >= 0:
+        e = entries[hi]
+        if e.run_length == 0:
+            return e
+        if tile_id - e.tile_id < e.run_length:
+            return e
+    return None
+
+
+class PMTilesReader:
+    """Random access to a PMTiles v3 archive on disk."""
+
+    def __init__(self, path: str | os.PathLike[str]) -> None:
+        self.path = os.fspath(path)
+        self._fh = open(self.path, "rb")  # noqa: SIM115 - closed in close()
+        self.header = Header.unpack(self._read(0, HEADER_SIZE))
+        self._root = self._dir(self.header.root_offset, self.header.root_length)
+        self._leaves: OrderedDict[tuple[int, int], list[Entry]] = OrderedDict()
+
+    def __enter__(self) -> PMTilesReader:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._fh.close()
+
+    def _read(self, offset: int, length: int) -> bytes:
+        self._fh.seek(offset)
+        data = self._fh.read(length)
+        if len(data) != length:
+            raise ValueError("unexpected end of PMTiles archive")
+        return data
+
+    def _dir(self, offset: int, length: int) -> list[Entry]:
+        return deserialize_directory(
+            _decompress(self._read(offset, length), self.header.internal_compression)
+        )
+
+    def _leaf(self, offset: int, length: int) -> list[Entry]:
+        key = (offset, length)
+        entries = self._leaves.get(key)
+        if entries is None:
+            entries = self._dir(self.header.leaf_offset + offset, length)
+            self._leaves[key] = entries
+            if len(self._leaves) > 64:
+                self._leaves.popitem(last=False)
+        else:
+            self._leaves.move_to_end(key)
+        return entries
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        h = self.header
+        if not h.metadata_length:
+            return {}
+        raw = _decompress(self._read(h.metadata_offset, h.metadata_length), h.internal_compression)
+        result: dict[str, Any] = json.loads(raw.decode("utf-8"))
+        return result
+
+    def get(self, tile: Tile | tuple[int, int, int]) -> bytes | None:
+        """Tile bytes, or ``None`` when the tile is not in the archive."""
+        z, x, y = tile
+        n = 1 << z
+        if not (0 <= x < n and 0 <= y < n):
+            return None
+        tid = zxy_to_tileid(z, x, y)
+        entries = self._root
+        for _ in range(4):  # the spec allows at most 3 directory levels below the root
+            e = _find(entries, tid)
+            if e is None:
+                return None
+            if e.run_length == 0:
+                entries = self._leaf(e.offset, e.length)
+                continue
+            return self._read(self.header.data_offset + e.offset, e.length)
+        raise ValueError("directory nesting is too deep")
+
+    def has(self, tile: Tile) -> bool:
+        return self.get(tile) is not None
+
+    def entries(self) -> Iterator[Entry]:
+        """All tile entries, leaf directories expanded, in tile-ID order."""
+
+        def walk(entries: list[Entry], depth: int) -> Iterator[Entry]:
+            for e in entries:
+                if e.run_length == 0:
+                    if depth > 3:
+                        raise ValueError("directory nesting is too deep")
+                    yield from walk(
+                        self._dir(self.header.leaf_offset + e.offset, e.length), depth + 1
+                    )
+                else:
+                    yield e
+
+        yield from walk(self._root, 1)
+
+    def tiles(self) -> Iterator[Tile]:
+        """Every addressed tile (runs expanded)."""
+        for e in self.entries():
+            for i in range(e.run_length):
+                yield tileid_to_zxy(e.tile_id + i)