tilefetch 0.1.0__py3-none-any.whl

tilefetch/storage.py

@@ -0,0 +1,338 @@
+"""Tile stores: MBTiles, PMTiles and plain ``z/x/y`` folders.
+
+Every store offers the same small interface used by the downloader:
+``has(tile)``, ``get(tile)``, ``put(tile, data)``, ``set_metadata(dict)``,
+``tiles()`` and ``close()``. Stores opened for writing keep what is already
+there, which is how interrupted downloads resume.
+"""
+
+from __future__ import annotations
+
+import os
+import sqlite3
+from typing import Any, Iterator, Protocol
+
+from .pmtiles import PMTilesReader, PMTilesWriter, tile_type_of, tileid_to_zxy, zxy_to_tileid
+from .tiles import Tile, flip_y
+
+__all__ = [
+    "TileStore",
+    "MBTilesStore",
+    "DirectoryStore",
+    "PMTilesStore",
+    "open_store",
+    "open_reader",
+    "store_kind",
+    "image_format",
+]
+
+_EXT_BY_TYPE = {2: "png", 3: "jpg", 4: "webp", 5: "avif"}
+
+
+def image_format(data: bytes) -> str | None:
+    """``"png"``, ``"jpg"``, ``"webp"`` or ``"avif"`` from magic bytes (else ``None``)."""
+    return _EXT_BY_TYPE.get(tile_type_of(data))
+
+
+class TileStore(Protocol):
+    def has(self, tile: Tile) -> bool: ...
+
+    def get(self, tile: Tile) -> bytes | None: ...
+
+    def put(self, tile: Tile, data: bytes) -> None: ...
+
+    def set_metadata(self, metadata: dict[str, Any]) -> None: ...
+
+    def tiles(self) -> Iterator[Tile]: ...
+
+    def close(self) -> None: ...
+
+
+def _meta_str(value: Any) -> str:
+    if isinstance(value, (list, tuple)):
+        return ",".join(_meta_str(v) for v in value)
+    if isinstance(value, float):
+        return repr(round(value, 7))
+    return str(value)
+
+
+class MBTilesStore:
+    """MBTiles 1.3 (SQLite) store.
+
+    Rows are stored in TMS order (``tile_row = 2**z - 1 - y``) as the spec
+    requires; the public API always speaks XYZ.
+    """
+
+    def __init__(self, path: str | os.PathLike[str], mode: str = "a") -> None:
+        self.path = os.fspath(path)
+        if mode == "r" and not os.path.exists(self.path):
+            raise FileNotFoundError(self.path)
+        if mode == "w" and os.path.exists(self.path):
+            os.remove(self.path)
+        self.readonly = mode == "r"
+        self._db = sqlite3.connect(self.path)
+        if not self.readonly:
+            self._db.executescript(
+                """
+                CREATE TABLE IF NOT EXISTS metadata (name TEXT, value TEXT);
+                CREATE UNIQUE INDEX IF NOT EXISTS metadata_name ON metadata (name);
+                CREATE TABLE IF NOT EXISTS tiles (
+                    zoom_level INTEGER, tile_column INTEGER, tile_row INTEGER, tile_data BLOB);
+                CREATE UNIQUE INDEX IF NOT EXISTS tile_index
+                    ON tiles (zoom_level, tile_column, tile_row);
+                """
+            )
+            self._db.commit()
+        self._pending = 0
+
+    def __enter__(self) -> MBTilesStore:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def has(self, tile: Tile) -> bool:
+        return self.get(tile) is not None
+
+    def get(self, tile: Tile) -> bytes | None:
+        z, x, y = tile
+        row = self._db.execute(
+            "SELECT tile_data FROM tiles WHERE zoom_level=? AND tile_column=? AND tile_row=?",
+            (z, x, flip_y(z, y)),
+        ).fetchone()
+        return bytes(row[0]) if row else None
+
+    def put(self, tile: Tile, data: bytes) -> None:
+        z, x, y = tile
+        self._db.execute(
+            "INSERT OR REPLACE INTO tiles (zoom_level, tile_column, tile_row, tile_data) "
+            "VALUES (?, ?, ?, ?)",
+            (z, x, flip_y(z, y), sqlite3.Binary(data)),
+        )
+        self._pending += 1
+        if self._pending >= 256:
+            self._db.commit()
+            self._pending = 0
+
+    @property
+    def metadata(self) -> dict[str, str]:
+        return {str(k): str(v) for k, v in self._db.execute("SELECT name, value FROM metadata")}
+
+    def set_metadata(self, metadata: dict[str, Any]) -> None:
+        for k, v in metadata.items():
+            self._db.execute(
+                "INSERT OR REPLACE INTO metadata (name, value) VALUES (?, ?)", (k, _meta_str(v))
+            )
+        self._db.commit()
+
+    def tiles(self) -> Iterator[Tile]:
+        rows = self._db.execute(
+            "SELECT zoom_level, tile_column, tile_row FROM tiles "
+            "ORDER BY zoom_level, tile_column, tile_row"
+        ).fetchall()
+        for z, x, row in rows:
+            yield Tile(int(z), int(x), flip_y(int(z), int(row)))
+
+    def __len__(self) -> int:
+        return int(self._db.execute("SELECT COUNT(*) FROM tiles").fetchone()[0])
+
+    def close(self) -> None:
+        try:
+            self._db.commit()
+        finally:
+            self._db.close()
+
+
+class DirectoryStore:
+    """A folder of ``{z}/{x}/{y}.{ext}`` files (XYZ rows, extension from content)."""
+
+    _EXTS = ("png", "jpg", "jpeg", "webp", "avif", "bin")
+
+    def __init__(self, root: str | os.PathLike[str], mode: str = "a") -> None:
+        self.root = os.fspath(root)
+        if mode == "r" and not os.path.isdir(self.root):
+            raise FileNotFoundError(self.root)
+        if mode != "r":
+            os.makedirs(self.root, exist_ok=True)
+
+    def __enter__(self) -> DirectoryStore:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def _find(self, tile: Tile) -> str | None:
+        z, x, y = tile
+        base = os.path.join(self.root, str(z), str(x), str(y))
+        for ext in self._EXTS:
+            p = f"{base}.{ext}"
+            if os.path.exists(p):
+                return p
+        return None
+
+    def has(self, tile: Tile) -> bool:
+        return self._find(tile) is not None
+
+    def get(self, tile: Tile) -> bytes | None:
+        p = self._find(tile)
+        if p is None:
+            return None
+        with open(p, "rb") as fh:
+            return fh.read()
+
+    def put(self, tile: Tile, data: bytes) -> None:
+        z, x, y = tile
+        folder = os.path.join(self.root, str(z), str(x))
+        os.makedirs(folder, exist_ok=True)
+        path = os.path.join(folder, f"{y}.{image_format(data) or 'bin'}")
+        tmp = path + ".tmp"
+        with open(tmp, "wb") as fh:
+            fh.write(data)
+        os.replace(tmp, path)
+
+    def set_metadata(self, metadata: dict[str, Any]) -> None:
+        import json
+
+        with open(os.path.join(self.root, "metadata.json"), "w", encoding="utf-8") as fh:
+            json.dump(metadata, fh, indent=2, ensure_ascii=False)
+
+    def tiles(self) -> Iterator[Tile]:
+        if not os.path.isdir(self.root):
+            return
+        for zd in sorted(os.listdir(self.root), key=lambda s: (len(s), s)):
+            if not zd.isdigit():
+                continue
+            zpath = os.path.join(self.root, zd)
+            for xd in sorted(os.listdir(zpath), key=lambda s: (len(s), s)):
+                if not xd.isdigit():
+                    continue
+                for fn in sorted(os.listdir(os.path.join(zpath, xd))):
+                    stem, _, ext = fn.partition(".")
+                    if stem.isdigit() and ext in self._EXTS:
+                        yield Tile(int(zd), int(xd), int(stem))
+
+    def close(self) -> None:
+        pass
+
+
+class PMTilesStore:
+    """Write-side PMTiles store (stages tiles, writes the archive on close)."""
+
+    def __init__(self, path: str | os.PathLike[str], mode: str = "a") -> None:
+        self.path = os.fspath(path)
+        if mode == "w":
+            for p in (self.path, self.path + ".part"):
+                if os.path.exists(p):
+                    os.remove(p)
+        part = self.path + ".part"
+        if mode == "a" and os.path.exists(self.path) and not os.path.exists(part):
+            # Resume on top of a finished archive: re-stage its tiles.
+            with PMTilesReader(self.path) as r:
+                meta = r.metadata
+                staged = [(t, r.get(t)) for t in r.tiles()]
+            self._writer = PMTilesWriter(self.path, meta)
+            for t, d in staged:
+                if d is not None:
+                    self._writer.put(t, d)
+        else:
+            self._writer = PMTilesWriter(self.path)
+
+    def __enter__(self) -> PMTilesStore:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        if exc[0] is None:
+            self.close()
+        else:
+            self.abort()
+
+    def has(self, tile: Tile) -> bool:
+        return self._writer.has(tile)
+
+    def get(self, tile: Tile) -> bytes | None:
+        row = self._writer._db.execute(
+            "SELECT data FROM tiles WHERE tile_id = ?",
+            (zxy_to_tileid(*tile),),
+        ).fetchone()
+        return bytes(row[0]) if row else None
+
+    def put(self, tile: Tile, data: bytes) -> None:
+        self._writer.put(tile, data)
+
+    def set_metadata(self, metadata: dict[str, Any]) -> None:
+        self._writer.set_metadata(metadata)
+
+    def tiles(self) -> Iterator[Tile]:
+        for (tid,) in self._writer._db.execute(
+            "SELECT tile_id FROM tiles ORDER BY tile_id"
+        ).fetchall():
+            yield tileid_to_zxy(int(tid))
+
+    def abort(self) -> None:
+        self._writer.abort()
+
+    def close(self) -> None:
+        self._writer.close()
+
+
+class _PMTilesReadStore:
+    """Read-only adapter so a PMTiles archive can feed :func:`tilefetch.mosaic`."""
+
+    def __init__(self, path: str | os.PathLike[str]) -> None:
+        self.reader = PMTilesReader(path)
+
+    def has(self, tile: Tile) -> bool:
+        return self.reader.has(tile)
+
+    def get(self, tile: Tile) -> bytes | None:
+        return self.reader.get(tile)
+
+    def put(self, tile: Tile, data: bytes) -> None:
+        raise OSError("PMTiles archives are opened read-only")
+
+    def set_metadata(self, metadata: dict[str, Any]) -> None:
+        raise OSError("PMTiles archives are opened read-only")
+
+    def tiles(self) -> Iterator[Tile]:
+        return self.reader.tiles()
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        return self.reader.metadata
+
+    def close(self) -> None:
+        self.reader.close()
+
+
+def store_kind(path: str | os.PathLike[str]) -> str:
+    """Classify an output path: ``mbtiles``, ``pmtiles``, ``image`` or ``dir``."""
+    p = os.fspath(path).lower()
+    if p.endswith(".mbtiles"):
+        return "mbtiles"
+    if p.endswith(".pmtiles"):
+        return "pmtiles"
+    if p.endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff", ".webp")):
+        return "image"
+    return "dir"
+
+
+def open_store(path: str | os.PathLike[str], mode: str = "a") -> TileStore:
+    """Open a writable store chosen by the path's extension (folder otherwise)."""
+    kind = store_kind(path)
+    if kind == "mbtiles":
+        return MBTilesStore(path, mode)
+    if kind == "pmtiles":
+        return PMTilesStore(path, mode)
+    if kind == "image":
+        raise ValueError("image outputs are produced by tilefetch.mosaic, not a tile store")
+    return DirectoryStore(path, mode)
+
+
+def open_reader(path: str | os.PathLike[str]) -> TileStore:
+    """Open an existing MBTiles file, PMTiles archive or tile folder for reading."""
+    kind = store_kind(path)
+    if kind == "mbtiles":
+        return MBTilesStore(path, "r")
+    if kind == "pmtiles":
+        return _PMTilesReadStore(path)
+    return DirectoryStore(path, "r")