Add tms parameter to cli for MosaicJSON (#233)

* start of better support for tms in cli * updated some tests * rename of tms parameter and adding to function def * Update test_cli.py * fix * optional tms * added stac asset for test in progress * should be working but seeing very odd behavior locally, may need to restart computer... * skip mars mosaic create if pyproj is too old * reformat --------- Co-authored-by: vincentsarago <[email protected]>
developmentseed · Oct 4, 2024 · b862d48 · b862d48
1 parent 45fd81e
commit b862d48
Show file tree

Hide file tree

Showing 8 changed files with 242 additions and 22 deletions.
diff --git a/cogeo_mosaic/backends/sqlite.py b/cogeo_mosaic/backends/sqlite.py
@@ -203,7 +203,7 @@ def write(self, overwrite: bool = False):
                         :layers
                     );
                 """,
-                self.mosaic_def.model_dump(exclude={"tiles"}),
+                self.mosaic_def.model_dump(exclude={"tiles"}, mode="json"),
             )
 
             self.db.executemany(
@@ -271,7 +271,7 @@ def update(
                         layers = :layers
                     WHERE name=:name
                 """,
-                self.mosaic_def.model_dump(exclude={"tiles"}),
+                self.mosaic_def.model_dump(exclude={"tiles"}, mode="json"),
             )
 
             if add_first:

diff --git a/cogeo_mosaic/mosaic.py b/cogeo_mosaic/mosaic.py
@@ -247,13 +247,15 @@ def from_urls(
         minzoom: Optional[int] = None,
         maxzoom: Optional[int] = None,
         max_threads: int = 20,
+        tilematrixset: Optional[morecantile.TileMatrixSet] = None,
         quiet: bool = True,
         **kwargs,
     ):
         """Create mosaicjson from COG urls.
 
         Attributes:
             urls (list): List of COGs.
+            tilematrixset: (morecantile.TileMatrixSet), optional (default: "WebMercatorQuad")
             minzoom (int): Force mosaic min-zoom.
             maxzoom (int): Force mosaic max-zoom.
             max_threads (int): Max threads to use (default: 20).
@@ -271,7 +273,9 @@ def from_urls(
             >>> MosaicJSON.from_urls(["1.tif", "2.tif"])
 
         """
-        features = get_footprints(urls, max_threads=max_threads, quiet=quiet)
+        features = get_footprints(
+            urls, max_threads=max_threads, tms=tilematrixset, quiet=quiet
+        )
 
         if minzoom is None:
             data_minzoom = {feat["properties"]["minzoom"] for feat in features}
@@ -297,7 +301,12 @@ def from_urls(
             raise MultipleDataTypeError("Dataset should have the same data type")
 
         return cls._create_mosaic(
-            features, minzoom=minzoom, maxzoom=maxzoom, quiet=quiet, **kwargs
+            features,
+            minzoom=minzoom,
+            maxzoom=maxzoom,
+            tilematrixset=tilematrixset,
+            quiet=quiet,
+            **kwargs,
         )
 
     @classmethod

diff --git a/cogeo_mosaic/scripts/cli.py b/cogeo_mosaic/scripts/cli.py
@@ -4,6 +4,7 @@
 import multiprocessing
 import os
 import sys
+from typing import Optional
 
 import click
 import cligj
@@ -20,7 +21,7 @@
 else:
     from importlib.metadata import entry_points
 
-tms = morecantile.tms.get("WebMercatorQuad")
+default_tms = morecantile.tms.get("WebMercatorQuad")
 
 
 @with_plugins(entry_points(group="cogeo_mosaic.plugins"))
@@ -59,6 +60,11 @@ def cogeo_cli():
     default=lambda: os.environ.get("MAX_THREADS", multiprocessing.cpu_count() * 5),
     help="threads",
 )
+@click.option(
+    "--tms",
+    help="Path to TileMatrixSet JSON file or valid TMS id.",
+    type=str,
+)
 @click.option("--name", type=str, help="Mosaic name")
 @click.option("--description", type=str, help="Mosaic description")
 @click.option("--attribution", type=str, help="Image attibution")
@@ -78,12 +84,22 @@ def create(
     min_tile_cover,
     tile_cover_sort,
     threads,
+    tms,
     name,
     description,
     attribution,
     quiet,
 ):
     """Create mosaic definition file."""
+    tilematrixset: Optional[morecantile.TileMatrixSet] = None
+    if tms:
+        if tms.endswith(".json"):
+            with open(tms, "r") as f:
+                tilematrixset = morecantile.TileMatrixSet(**json.load(f))
+
+        else:
+            tilematrixset = morecantile.tms.get(tms)
+
     input_files = [file.strip() for file in input_files if file.strip()]
     mosaicjson = MosaicJSON.from_urls(
         input_files,
@@ -92,6 +108,7 @@ def create(
         quadkey_zoom=quadkey_zoom,
         minimum_tile_cover=min_tile_cover,
         tile_cover_sort=tile_cover_sort,
+        tilematrixset=tilematrixset,
         max_threads=threads,
         quiet=quiet,
     )
@@ -104,7 +121,11 @@ def create(
         mosaicjson.attribution = attribution
 
     if output:
-        with MosaicBackend(output, mosaic_def=mosaicjson) as mosaic:
+        with MosaicBackend(
+            output,
+            mosaic_def=mosaicjson,
+            tms=tilematrixset or default_tms,
+        ) as mosaic:
             mosaic.write(overwrite=True)
     else:
         click.echo(mosaicjson.model_dump_json(exclude_none=True))
@@ -115,11 +136,26 @@ def create(
 @click.option(
     "--url", type=str, required=True, help="URL to which the mosaic should be uploaded."
 )
-def upload(file, url):
+@click.option(
+    "--tms",
+    help="Path to TileMatrixSet JSON file or valid TMS id.",
+    type=str,
+)
+def upload(file, url, tms):
     """Upload mosaic definition file."""
-    mosaicjson = json.load(file)
+    tilematrixset: Optional[morecantile.TileMatrixSet] = None
+    if tms:
+        if tms.endswith(".json"):
+            with open(tms, "r") as f:
+                tilematrixset = morecantile.TileMatrixSet(**json.load(f))
 
-    with MosaicBackend(url, mosaic_def=mosaicjson) as mosaic:
+        else:
+            tilematrixset = morecantile.tms.get(tms)
+
+    mosaicjson = MosaicJSON.model_validate(json.load(file))
+
+    tilematrixset = tilematrixset or mosaicjson.tilematrixset or default_tms
+    with MosaicBackend(url, mosaic_def=mosaicjson, tms=tilematrixset) as mosaic:
         mosaic.write(overwrite=True)
 
 
@@ -140,6 +176,11 @@ def upload(file, url):
 @click.option(
     "--tile-cover-sort", help="Sort files by covering %", is_flag=True, default=False
 )
+@click.option(
+    "--tms",
+    help="Path to TileMatrixSet JSON file or valid TMS id.",
+    type=str,
+)
 @click.option("--name", type=str, help="Mosaic name")
 @click.option("--description", type=str, help="Mosaic description")
 @click.option("--attribution", type=str, help="Image attibution")
@@ -159,12 +200,22 @@ def create_from_features(
     quadkey_zoom,
     min_tile_cover,
     tile_cover_sort,
+    tms,
     name,
     description,
     attribution,
     quiet,
 ):
     """Create mosaic definition file."""
+    tilematrixset: Optional[morecantile.TileMatrixSet] = None
+    if tms:
+        if tms.endswith(".json"):
+            with open(tms, "r") as f:
+                tilematrixset = morecantile.TileMatrixSet(**json.load(f))
+
+        else:
+            tilematrixset = morecantile.tms.get(tms)
+
     mosaicjson = MosaicJSON.from_features(
         list(features),
         minzoom,
@@ -173,6 +224,7 @@ def create_from_features(
         accessor=lambda feature: feature["properties"][property],
         minimum_tile_cover=min_tile_cover,
         tile_cover_sort=tile_cover_sort,
+        tilematrixset=tilematrixset,
         quiet=quiet,
     )
 
@@ -184,7 +236,11 @@ def create_from_features(
         mosaicjson.attribution = attribution
 
     if output:
-        with MosaicBackend(output, mosaic_def=mosaicjson) as mosaic:
+        with MosaicBackend(
+            output,
+            mosaic_def=mosaicjson,
+            tms=tilematrixset or default_tms,
+        ) as mosaic:
             mosaic.write(overwrite=True)
     else:
         click.echo(mosaicjson.model_dump_json(exclude_none=True))
@@ -200,6 +256,11 @@ def create_from_features(
     is_flag=True,
     default=True,
 )
+@click.option(
+    "--tms",
+    help="Path to TileMatrixSet JSON file or valid TMS id.",
+    type=str,
+)
 @click.option(
     "--threads",
     type=int,
@@ -213,11 +274,20 @@ def create_from_features(
     is_flag=True,
     default=False,
 )
-def update(input_files, input_mosaic, min_tile_cover, add_first, threads, quiet):
+def update(input_files, input_mosaic, min_tile_cover, tms, add_first, threads, quiet):
     """Update mosaic definition file."""
+    tilematrixset: Optional[morecantile.TileMatrixSet] = None
+    if tms:
+        if tms.endswith(".json"):
+            with open(tms, "r") as f:
+                tilematrixset = morecantile.TileMatrixSet(**json.load(f))
+
+        else:
+            tilematrixset = morecantile.tms.get(tms)
+
     input_files = input_files.read().splitlines()
-    features = get_footprints(input_files, max_threads=threads)
-    with MosaicBackend(input_mosaic) as mosaic:
+    features = get_footprints(input_files, tms=tilematrixset, max_threads=threads)
+    with MosaicBackend(input_mosaic, tms=tilematrixset or default_tms) as mosaic:
         mosaic.update(
             features,
             add_first=add_first,
@@ -235,18 +305,34 @@ def update(input_files, input_mosaic, min_tile_cover, add_first, threads, quiet)
     default=lambda: os.environ.get("MAX_THREADS", multiprocessing.cpu_count() * 5),
     help="threads",
 )
+@click.option(
+    "--tms",
+    help="Path to TileMatrixSet JSON file or valid TMS id.",
+    type=str,
+)
 @click.option(
     "--quiet",
     "-q",
     help="Remove progressbar and other non-error output.",
     is_flag=True,
     default=False,
 )
-def footprint(input_files, output, threads, quiet):
+def footprint(input_files, output, threads, tms, quiet):
     """Create mosaic definition file."""
+    tilematrixset: Optional[morecantile.TileMatrixSet] = None
+    if tms:
+        if tms.endswith(".json"):
+            with open(tms, "r") as f:
+                tilematrixset = morecantile.TileMatrixSet(**json.load(f))
+
+        else:
+            tilematrixset = morecantile.tms.get(tms)
+
     input_files = input_files.read().splitlines()
     foot = {
-        "features": get_footprints(input_files, max_threads=threads, quiet=quiet),
+        "features": get_footprints(
+            input_files, max_threads=threads, tms=tilematrixset, quiet=quiet
+        ),
         "type": "FeatureCollection",
     }
 
@@ -266,9 +352,26 @@ def footprint(input_files, output, threads, quiet):
     is_flag=True,
     help="Print as JSON.",
 )
-def info(input, to_json):
+@click.option(
+    "--tms",
+    help="Path to TileMatrixSet JSON file or valid TMS id.",
+    type=str,
+)
+def info(input, to_json, tms):
     """Return info about the mosaic."""
-    with MosaicBackend(input) as mosaic:
+    tilematrixset: Optional[morecantile.TileMatrixSet] = None
+    if tms:
+        if tms.endswith(".json"):
+            with open(tms, "r") as f:
+                tilematrixset = morecantile.TileMatrixSet(**json.load(f))
+
+        else:
+            tilematrixset = morecantile.tms.get(tms)
+
+    with MosaicBackend(
+        input,
+        tms=tilematrixset or default_tms,
+    ) as mosaic:
         _info = {
             "Path": input,
             "Backend": mosaic._backend_name,
@@ -285,7 +388,7 @@ def info(input, to_json):
         }
 
         geo = {
-            "TileMatrixSet": "WebMercatorQuad",
+            "TileMatrixSet": mosaic.tms.id,
             "BoundingBox": mosaic.mosaic_def.bounds,
             "Center": mosaic.mosaic_def.center,
             "Min Zoom": mosaic.mosaic_def.minzoom,
@@ -352,8 +455,8 @@ def to_geojson(input, collect):
     features = []
     with MosaicBackend(input) as mosaic:
         for qk, assets in mosaic.mosaic_def.tiles.items():
-            tile = tms.quadkey_to_tile(qk)
-            west, south, east, north = tms.bounds(tile)
+            tile = mosaic.tms.quadkey_to_tile(qk)
+            west, south, east, north = mosaic.tms.bounds(tile)
 
             geom = {
                 "type": "Polygon",

diff --git a/cogeo_mosaic/utils.py b/cogeo_mosaic/utils.py
@@ -5,7 +5,7 @@
 import sys
 from concurrent import futures
 from contextlib import ExitStack
-from typing import Dict, List, Sequence, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple
 
 import click
 import morecantile
@@ -78,7 +78,7 @@ def get_dataset_info(
 
 def get_footprints(
     dataset_list: Sequence[str],
-    tms: morecantile.TileMatrixSet = WEB_MERCATOR_TMS,
+    tms: Optional[morecantile.TileMatrixSet] = None,
     max_threads: int = 20,
     quiet: bool = True,
 ) -> List:
@@ -100,6 +100,8 @@ def get_footprints(
         tuple of footprint feature.
 
     """
+    tms = tms or WEB_MERCATOR_TMS
+
     with ExitStack() as ctx:
         fout = ctx.enter_context(open(os.devnull, "w")) if quiet else sys.stderr
         with futures.ThreadPoolExecutor(max_workers=max_threads) as executor: