diff --git a/elevation/__init__.py b/elevation/__init__.py index a7a30a7..a8fafc7 100644 --- a/elevation/__init__.py +++ b/elevation/__init__.py @@ -21,8 +21,8 @@ # Local copy or not installed with setuptools __version__ = "999" +from .cache import CACHE_DIR, resolve_cache_dir from .datasource import ( - CACHE_DIR, DEFAULT_GDAL_OPTIONS, DEFAULT_OUTPUT, DEFAULT_PRODUCT, @@ -35,7 +35,6 @@ dataset, distclean, info, - resolve_cache_dir, seed, ) from .spatial import TOOLS, selfcheck diff --git a/elevation/__main__.py b/elevation/__main__.py index ec4fac6..8b46bd4 100644 --- a/elevation/__main__.py +++ b/elevation/__main__.py @@ -121,8 +121,14 @@ def seed( "--bounds", help="Output bounds: left bottom right top.", ), + margin: str = typer.Option( + elevation.MARGIN, + "-m", + "--margin", + help="Decimal degree margin added to the bounds. Use '%' for percent margin.", + ), ) -> None: - elevation.seed(**ctx.obj, bounds=bounds) + elevation.seed(**ctx.obj, bounds=bounds, margin=margin) @app.command(short_help="Clip the DEM to given bounds.") diff --git a/elevation/cache.py b/elevation/cache.py new file mode 100644 index 0000000..4b0f3b7 --- /dev/null +++ b/elevation/cache.py @@ -0,0 +1,118 @@ +# +# Copyright (c) 2016-2026 B-Open Solutions srl - https://bopen.eu +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shutil +from collections.abc import Callable, Generator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import appdirs +import fasteners + +from . import spatial + +CACHE_DIR: str = appdirs.user_cache_dir("elevation", "bopen") +FOLDER_LOCKFILE_NAME = ".folder_lock" +Tile = tuple[tuple[int, int], str] + + +def resolve_cache_dir(cache_dir: str | Path | None) -> Path: + """Return the DEM cache folder to use, as an absolute path. + + The ``cache_dir`` argument takes precedence over the ``EIO_CACHE_DIR`` environment + variable, that takes precedence over the ``CACHE_DIR`` default. + """ + if cache_dir is None: + cache_dir = os.environ.get("EIO_CACHE_DIR") or CACHE_DIR + return Path(cache_dir).resolve() + + +def ensure_setup(root: Path) -> None: + """Create the product folder and its ``cache`` subfolder. + + The ``spool`` folder is created on demand by the tile download. + """ + with fasteners.InterProcessLock(root / FOLDER_LOCKFILE_NAME): + for path in (root, root / "cache"): + path.mkdir(parents=True, exist_ok=True) + + +@contextmanager +def lock_tiles(datasource_root: Path, tile_names: list[str]) -> Generator[None]: + locks = [] + for tile_name in tile_names: + lockfile = datasource_root / "cache" / f"{tile_name}.lock" + locks.append(fasteners.InterProcessLock(lockfile)) + + for lock in locks: + lock.acquire(blocking=True) + + yield + + for lock in locks: + lock.release() + + +def ensure_tiles( + root: Path, + tiles: list[Tile], + prepare_tile: Callable[..., tuple[str, Path | None]], + gdal_options: str = spatial.INT_TILE_GDAL_OPTIONS, + **kwargs: Any, +) -> None: + """Fetch and cache *tiles*, skipping the tiles already in the cache. + + A tile is a ``(name, window)`` pair: a tile with a window is read in place + from ``datasource_url``, a tile without one is downloaded whole from its own + URL and goes through the spool. + """ + for (ilon, ilat), tile_name in tiles: + cached = root / "cache" / tile_name + if cached.exists() and cached.stat().st_size > 0: + continue + + # prepare the data if GDAL cannot download it / read it as it is + source, spooled = prepare_tile( + tile_name=tile_name, spool=root / "spool", ilat=ilat, ilon=ilon, **kwargs + ) + + # convert the data to the internal cache format + ready = root / "spool/ready" / tile_name + spatial.call_gdal_translate(source, ready, options=gdal_options) + if spooled is not None: + spooled.unlink(missing_ok=True) + + # finally move the data inside the cache. The move is atomic in most cases + cached.parent.mkdir(parents=True, exist_ok=True) + shutil.move(ready, cached) + + +@contextmanager +def lock_vrt(datasource_root: Path, product: str) -> Generator[None]: + with fasteners.InterProcessLock(datasource_root / f"{product}.vrt.lock"): + yield + + +def build_vrt(root: Path, product: str) -> list[str]: + """Build the ``.vrt`` mosaic over the non empty cache tiles.""" + tiles = [] + for tile in (root / "cache").rglob("*.tif"): + if tile.stat().st_size > 0: + tiles.append(str(tile)) + options = "-q -overwrite" + cmd = spatial.call_gdalbuildvrt(sorted(tiles), root / f"{product}.vrt", options) + return cmd diff --git a/elevation/datasource.py b/elevation/datasource.py index 51083a9..0f528c1 100644 --- a/elevation/datasource.py +++ b/elevation/datasource.py @@ -14,39 +14,16 @@ # limitations under the License. import math -import os import shutil -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Iterator from importlib import resources from pathlib import Path from typing import Any, NotRequired, TypedDict -import appdirs - -from . import spatial, util - -__all__ = [ - "CACHE_DIR", - "DEFAULT_GDAL_OPTIONS", - "DEFAULT_OUTPUT", - "DEFAULT_PRODUCT", - "MARGIN", - "PRODUCTS", - "RETIRED_PRODUCTS", - "ProductRetiredError", - "clean", - "clip", - "dataset", - "distclean", - "info", - "resolve_cache_dir", - "seed", -] - -CACHE_DIR: str = appdirs.user_cache_dir("elevation", "bopen") +from . import cache, spatial + DEFAULT_OUTPUT = "out.tif" DEFAULT_GDAL_OPTIONS = "-co TILED=YES -co COMPRESS=DEFLATE -co ZLEVEL=9 -co PREDICTOR=2" -CACHE_EXT = ".tif" MARGIN = "0" # NOTE: @@ -58,17 +35,6 @@ CGIAR_L1_TILE_INDECES_TRANSFORM = (-185.0004166666667, 5.0, 65.0004166666667, -5.0) -def resolve_cache_dir(cache_dir: str | Path | None) -> Path: - """Return the DEM cache folder to use, as an absolute path. - - The ``cache_dir`` argument takes precedence over the ``EIO_CACHE_DIR`` environment - variable, that takes precedence over the ``CACHE_DIR`` default. - """ - if cache_dir is None: - cache_dir = os.environ.get("EIO_CACHE_DIR") or CACHE_DIR - return Path(cache_dir).resolve() - - def latlon_to_indeces( transform: tuple[float, float, float, float], lon: float, lat: float ) -> tuple[int, int]: @@ -78,16 +44,13 @@ def latlon_to_indeces( return ilon, ilat -Tile = tuple[tuple[int, int], str] - - def dted_l2_tiles( left: float, bottom: float, right: float, top: float, tile_name_template: str = "{slat}{slon}.tif", -) -> Iterator[Tile]: +) -> Iterator[cache.Tile]: ileft, itop = latlon_to_indeces(DTED_L2_TILE_INDECES_TRANSFORM, left, top) iright, ibottom = latlon_to_indeces(DTED_L2_TILE_INDECES_TRANSFORM, right, bottom) # special case often used *integer* top and right to avoid downloading unneeded tiles @@ -108,7 +71,7 @@ def cgiar_l1_tiles( right: float, top: float, tile_template: str = "srtm_{ilon:02d}_{ilat:02d}.tif", -) -> Iterator[Tile]: +) -> Iterator[cache.Tile]: ileft, itop = latlon_to_indeces(CGIAR_L1_TILE_INDECES_TRANSFORM, left, top) iright, ibottom = latlon_to_indeces(CGIAR_L1_TILE_INDECES_TRANSFORM, right, bottom) for ilon in range(ileft, iright + 1): @@ -123,7 +86,7 @@ def srtm_ellip_tiles( right: float, top: float, tile_name_template: str = "{slat}{slon}_wgs84.tif", -) -> Iterator[Tile]: +) -> Iterator[cache.Tile]: ileft, itop = latlon_to_indeces(DTED_L2_TILE_INDECES_TRANSFORM, left, top) iright, ibottom = latlon_to_indeces(DTED_L2_TILE_INDECES_TRANSFORM, right, bottom) # special case often used *integer* top and right to avoid downloading unneeded tiles @@ -167,7 +130,7 @@ def zarr_tiles( right: float, top: float, transform: tuple[float, float, float, float], -) -> Iterator[Tile]: +) -> Iterator[cache.Tile]: ileft, itop = latlon_to_indeces(transform, left, top) iright, ibottom = latlon_to_indeces(transform, right, bottom) for ilon in range(ileft, iright + 1): @@ -193,18 +156,11 @@ def prepare_tile_zarr( class DatasourceSpec(TypedDict): - # a local product has one URL per tile (``tiles``), a remote one is a - # single chunked source (``grid``): the key tells the two apart - cached_tiles: Callable[..., Iterator[Tile]] - # keyword arguments for ``cached_tiles``, e.g. the tile name template - # of a product that keeps its tiles in subfolders + """How a DEM product lists, prepares and caches its tiles.""" + + cached_tiles: Callable[..., Iterator[cache.Tile]] cached_tiles_kwargs: NotRequired[dict[str, Any]] - # prepare the tile for GDAL, downloading it or reading the window of the - # chunked source, next to the spool file to remove once it is cached prepare_tile: Callable[..., tuple[str, Path | None]] - # keyword arguments for ``prepare_tile``: the datasource URL and the path - # of the variable in the store, the source extension, the archive the - # provider serves it in, the chunk size prepare_tile_kwargs: dict[str, Any] tile_gdal_options: NotRequired[str] @@ -316,7 +272,7 @@ def tile_source( the ``compressed_ext`` when the source is compressed. The member is the file to read inside a ``.zip`` archive and ``None`` otherwise. """ - stem = tile_name.removesuffix(CACHE_EXT) + stem = tile_name.removesuffix(".tif") spool_name = f"{stem}{tile_ext}" remote = spool_name if compressed_ext is None else f"{stem}{compressed_ext}" member = Path(spool_name).name if compressed_ext == ".zip" else None @@ -352,83 +308,24 @@ def fetch_tile(source: str, destination: Path, *, member: str | None = None) -> temporary.unlink(missing_ok=True) -def ensure_tiles( - root: Path, - tiles: Sequence[Tile], - prepare_tile: Callable[..., tuple[str, Path | None]], - gdal_options: str = spatial.INT_TILE_GDAL_OPTIONS, - **kwargs: Any, -) -> None: - """Fetch and cache *tiles*, skipping the tiles already in the cache. - - A tile is a ``(name, window)`` pair: a tile with a window is read in place - from ``datasource_url``, a tile without one is downloaded whole from its own - URL and goes through the spool. - """ - for (ilon, ilat), tile_name in tiles: - cached = root / "cache" / tile_name - if cached.exists() and cached.stat().st_size > 0: - continue - - # prepare the data if GDAL cannot download it / read it as it is - source, spooled = prepare_tile( - tile_name=tile_name, spool=root / "spool", ilat=ilat, ilon=ilon, **kwargs - ) - - # convert the data to the internal cache format - ready = root / "spool/ready" / tile_name - spatial.call_gdal_translate(source, ready, options=gdal_options) - if spooled is not None: - spooled.unlink(missing_ok=True) - - # finally move the data inside the cache. The move is atomic in most cases - cached.parent.mkdir(parents=True, exist_ok=True) - shutil.move(ready, cached) - - -def build_vrt(root: Path, product: str) -> list[str]: - """Build the ``.vrt`` mosaic over the non empty cache tiles.""" - tiles = [] - for tile in (root / "cache").rglob("*.tif"): - if tile.stat().st_size > 0: - tiles.append(str(tile)) - options = "-q -overwrite" - cmd = spatial.call_gdalbuildvrt(sorted(tiles), root / f"{product}.vrt", options) - return cmd - - def ensure_setup( cache_dir: str | Path | None, product: str ) -> tuple[Path, DatasourceSpec]: if product in RETIRED_PRODUCTS: raise ProductRetiredError(RETIRED_PRODUCTS[product]) - datasource_root = resolve_cache_dir(cache_dir) / product + datasource_root = cache.resolve_cache_dir(cache_dir) / product spec = PRODUCTS_SPECS[product] - util.ensure_setup(datasource_root) + cache.ensure_setup(datasource_root) return datasource_root, spec -def do_clip( - path: Path, - bounds: tuple[float, float, float, float], - output: Path, - product: str, - gdal_options: str = DEFAULT_GDAL_OPTIONS, -) -> list[str]: - left, bottom, right, top = bounds - options = f"-q {gdal_options} -projwin {left} {top} {right} {bottom}" - source = str(path / f"{product}.vrt") - with util.lock_vrt(path, product): - cmd = spatial.call_gdal_translate(source, output, options=options) - return cmd - - def seed( cache_dir: str | Path | None = None, product: str = DEFAULT_PRODUCT, bounds: tuple[float, float, float, float] | None = None, + margin: str = MARGIN, max_download_tiles: int = 9, -) -> Path: +) -> tuple[Path, tuple[float, float, float, float]]: """Seed the DEM to given bounds. A remote product is not downloaded whole: only the chunks of the store that @@ -437,10 +334,13 @@ def seed( :param cache_dir: Root of the DEM cache folder. :param product: DEM product choice. :param bounds: Output bounds in 'left bottom right top' order. + :param margin: Decimal degree margin added to the bounds. Use '%' for percent margin. :param max_download_tiles: Maximum number of tiles to process. + :return: The datasource root and the bounds with the margin applied. """ if bounds is None: raise TypeError("bounds must be supplied") + bounds = build_bounds(bounds, margin=margin) datasource_root, spec = ensure_setup(cache_dir, product) cached_tiles = spec["cached_tiles"] cached_tiles_kwargs = spec.get("cached_tiles_kwargs", {}) @@ -454,8 +354,8 @@ def seed( prepare_tile = spec["prepare_tile"] prepare_tile_kwargs = spec.get("prepare_tile_kwargs", {}) - with util.lock_tiles(datasource_root, [name for _, name in tiles]): - ensure_tiles( + with cache.lock_tiles(datasource_root, [name for _, name in tiles]): + cache.ensure_tiles( datasource_root, tiles, prepare_tile=prepare_tile, @@ -463,10 +363,10 @@ def seed( **prepare_tile_kwargs, ) - with util.lock_vrt(datasource_root, product): - build_vrt(datasource_root, product) + with cache.lock_vrt(datasource_root, product): + cache.build_vrt(datasource_root, product) - return datasource_root + return datasource_root, bounds def build_bounds( @@ -506,9 +406,14 @@ def clip( :param gdal_options: GDAL creation options of the output file. """ output = Path(output).resolve() - bounds = build_bounds(bounds, margin=margin) - datasource_root = seed(cache_dir=cache_dir, product=product, bounds=bounds) - do_clip(datasource_root, bounds, output, product=product, gdal_options=gdal_options) + datasource_root, bounds = seed( + cache_dir=cache_dir, product=product, bounds=bounds, margin=margin + ) + left, bottom, right, top = bounds + options = f"-q {gdal_options} -projwin {left} {top} {right} {bottom}" + source = str(datasource_root / f"{product}.vrt") + with cache.lock_vrt(datasource_root, product): + spatial.call_gdal_translate(source, output, options=options) def dataset(dataset: str | None = None) -> str: diff --git a/elevation/util.py b/elevation/util.py deleted file mode 100644 index f0fa46d..0000000 --- a/elevation/util.py +++ /dev/null @@ -1,54 +0,0 @@ -# -# Copyright (c) 2016-2026 B-Open Solutions srl - https://bopen.eu -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from collections.abc import Generator, Iterable -from contextlib import contextmanager -from pathlib import Path - -import fasteners - -FOLDER_LOCKFILE_NAME = ".folder_lock" - - -@contextmanager -def lock_tiles(datasource_root: Path, tile_names: Iterable[str]) -> Generator[None]: - locks = [] - for tile_name in tile_names: - lockfile = datasource_root / "cache" / f"{tile_name}.lock" - locks.append(fasteners.InterProcessLock(lockfile)) - - for lock in locks: - lock.acquire(blocking=True) - - yield - - for lock in locks: - lock.release() - - -@contextmanager -def lock_vrt(datasource_root: Path, product: str) -> Generator[None]: - with fasteners.InterProcessLock(datasource_root / f"{product}.vrt.lock"): - yield - - -def ensure_setup(root: Path) -> None: - """Create the product folder and its ``cache`` subfolder. - - The ``spool`` folder is created on demand by the tile download. - """ - with fasteners.InterProcessLock(root / FOLDER_LOCKFILE_NAME): - for path in (root, root / "cache"): - path.mkdir(parents=True, exist_ok=True) diff --git a/tests/test_10_util.py b/tests/test_15_cache.py similarity index 79% rename from tests/test_10_util.py rename to tests/test_15_cache.py index 6213fff..7e4b7ca 100644 --- a/tests/test_10_util.py +++ b/tests/test_15_cache.py @@ -4,26 +4,26 @@ from pathlib import Path -from elevation import util +from elevation import cache def test_lock_tiles(tmp_path: Path) -> None: root = tmp_path / "root" - with util.lock_tiles(root, ["a.tiff"]): + with cache.lock_tiles(root, ["a.tiff"]): assert (root / "cache" / "a.tiff.lock").exists() def test_lock_vrt(tmp_path: Path) -> None: root = tmp_path / "root" - with util.lock_vrt(root, "SRTM1_GEOID"): + with cache.lock_vrt(root, "SRTM1_GEOID"): assert (root / "SRTM1_GEOID.vrt.lock").exists() def test_ensure_setup(tmp_path: Path) -> None: root = tmp_path / "root" - util.ensure_setup(root) + cache.ensure_setup(root) assert (root / "cache").is_dir() # the spool folder is created on demand by the tile download diff --git a/tests/test_20_datasource.py b/tests/test_20_datasource.py index 80686d1..cf3f6ac 100644 --- a/tests/test_20_datasource.py +++ b/tests/test_20_datasource.py @@ -9,7 +9,7 @@ from pytest_mock import MockerFixture import elevation -from elevation import datasource, spatial +from elevation import cache, datasource, spatial DATA_DIR = Path(__file__).parent / "data" REFERENCE = DATA_DIR / "reference.tif" @@ -178,7 +178,7 @@ def test_ensure_tiles(mocker: MockerFixture, tmp_path: Path) -> None: "elevation.spatial.call_gdal_translate", side_effect=write_ready_tile ) - datasource.ensure_tiles( + cache.ensure_tiles( tmp_path, [((12, 41), "N41E012.tif")], prepare_tile=spec["prepare_tile"], @@ -207,7 +207,7 @@ def test_ensure_tiles_skips_cached(mocker: MockerFixture, tmp_path: Path) -> Non mock_fetch = mocker.patch("elevation.datasource.fetch_tile") mock_write = mocker.patch("elevation.spatial.call_gdal_translate") - datasource.ensure_tiles( + cache.ensure_tiles( tmp_path, [((12, 41), "N41E012.tif")], prepare_tile=spec["prepare_tile"], @@ -231,7 +231,7 @@ def test_ensure_tiles_remote(mocker: MockerFixture, tmp_path: Path) -> None: ) assert tiles == [((96, 24), "24/96.tif")] - datasource.ensure_tiles( + cache.ensure_tiles( tmp_path, tiles, prepare_tile=spec["prepare_tile"], @@ -278,58 +278,6 @@ def test_fetch_tile_zip(tmp_path: Path) -> None: assert destination.read_bytes() == REFERENCE.read_bytes() -def test_do_clip(mocker: MockerFixture, tmp_path: Path) -> None: - bounds = (13.1, 43.1, 14.9, 44.9) - mock_check_call = mocker.patch("subprocess.check_call") - - cmd = datasource.do_clip( - path=tmp_path, bounds=bounds, output=Path("/out.tif"), product="SRTM3" - ) - - expected_cmd = [ - "gdal_translate", - "-q", - *datasource.DEFAULT_GDAL_OPTIONS.split(), - "-projwin", - "13.1", - "44.9", - "14.9", - "43.1", - str(tmp_path / "SRTM3.vrt"), - "/out.tif", - ] - assert cmd == expected_cmd - mock_check_call.assert_called_once_with(cmd) - - -def test_do_clip_gdal_options(mocker: MockerFixture, tmp_path: Path) -> None: - mock_check_call = mocker.patch("subprocess.check_call") - - cmd = datasource.do_clip( - path=tmp_path, - bounds=(1.0, 2.0, 3.0, 4.0), - output=Path("/out.tif"), - product="SRTM3", - gdal_options="-co COMPRESS=LZW", - ) - - expected_cmd = [ - "gdal_translate", - "-q", - "-co", - "COMPRESS=LZW", - "-projwin", - "1.0", - "4.0", - "3.0", - "2.0", - str(tmp_path / "SRTM3.vrt"), - "/out.tif", - ] - assert cmd == expected_cmd - mock_check_call.assert_called_once_with(cmd) - - def test_seed(mocker: MockerFixture, tmp_path: Path) -> None: root = tmp_path / "root" bounds = (13.1, 43.1, 13.9, 43.9) @@ -340,11 +288,12 @@ def test_seed(mocker: MockerFixture, tmp_path: Path) -> None: "elevation.spatial.call_gdal_translate", side_effect=write_ready_tile ) - datasource_root = datasource.seed( + datasource_root, seeded_bounds = datasource.seed( cache_dir=root, product="SRTM1_GEOID", bounds=bounds ) assert datasource_root == root / "SRTM1_GEOID" + assert seeded_bounds == bounds mock_fetch.assert_called_once_with( f"{spec['prepare_tile_kwargs']['datasource_url']}/N43E013.tif", datasource_root / "spool" / "N43E013.tif", @@ -372,7 +321,7 @@ def test_seed_remote(mocker: MockerFixture, tmp_path: Path) -> None: "elevation.spatial.call_gdal_translate", side_effect=write_ready_tile ) - datasource_root = datasource.seed( + datasource_root, _ = datasource.seed( cache_dir=root, product="GLO-30", bounds=(12.4, 41.8, 12.4 + 100 / 3600, 41.8 + 100 / 3600), @@ -467,7 +416,7 @@ def test_cache_dir(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: default = tmp_path / "default" override = tmp_path / "override" argument = tmp_path / "argument" - monkeypatch.setattr(datasource, "CACHE_DIR", default) + monkeypatch.setattr(cache, "CACHE_DIR", default) assert f"Product folder: {default.resolve() / 'MAPZEN'}" in datasource.info() diff --git a/tests/test_40_main.py b/tests/test_40_main.py index 5b3b71e..721542a 100644 --- a/tests/test_40_main.py +++ b/tests/test_40_main.py @@ -106,6 +106,18 @@ def test_eio_seed(mocker: MockerFixture, tmp_path: Path) -> None: assert mock_check_call.call_args[0][0][0] == "gdalbuildvrt" +def test_eio_seed_margin(mocker: MockerFixture, tmp_path: Path) -> None: + root = tmp_path / "root" + runner = typer.testing.CliRunner() + mock_seed = mocker.patch("elevation.seed") + result = runner.invoke( + __main__.app, + f"--cache_dir {root!s} seed --bounds 12.5 42 12.5 42 -m 1".split(), + ) + assert not result.exception + assert mock_seed.call_args.kwargs["margin"] == "1" + + def test_eio_clip(mocker: MockerFixture, tmp_path: Path) -> None: root = tmp_path / "root" runner = typer.testing.CliRunner()