feat: add existing tile list

ncx-co · Dec 21, 2022 · 8702315 · 8702315
1 parent 957fa53
commit 8702315
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/src/stactools/usda_cdl/tile.py b/src/stactools/usda_cdl/tile.py
@@ -47,14 +47,20 @@ def tile_zipfile(
     directory: Path,
     size: int = DEFAULT_WINDOW_SIZE,
     max_workers: int = DEFAULT_MAX_WORKERS,
+    existing_tiles: Optional[List[str]] = None,
 ) -> List[Path]:
     """Tiles an input GeoTIFF (wrapped in a zipfile)."""
     if infile.suffix != ".zip":
         raise ValueError(f"Infile should end in .zip: {infile}")
     zip_path = f"zip://{infile}!/{infile.stem}.tif"
     with rasterio.open(zip_path) as dataset:
         return _tile_dataset(
-            dataset, Metadata.from_href(infile.stem), directory, size, max_workers
+            dataset,
+            Metadata.from_href(infile.stem),
+            directory,
+            size,
+            max_workers,
+            existing_tiles or list(),
         )
 
 
@@ -63,11 +69,17 @@ def tile_geotiff(
     directory: Path,
     size: int = DEFAULT_WINDOW_SIZE,
     max_workers: int = DEFAULT_MAX_WORKERS,
+    existing_tiles: Optional[List[str]] = None,
 ) -> List[Path]:
     """Tiles an input GeoTIFF."""
     with rasterio.open(infile) as dataset:
         return _tile_dataset(
-            dataset, Metadata.from_href(str(infile)), directory, size, max_workers
+            dataset,
+            Metadata.from_href(str(infile)),
+            directory,
+            size,
+            max_workers,
+            existing_tiles or list(),
         )
 
 
@@ -77,11 +89,15 @@ def _tile_dataset(
     directory: Path,
     size: int,
     max_workers: int,
+    existing_tiles: List[str],
 ) -> List[Path]:
     windows = _create_windows(dataset, size)
     read_lock = threading.Lock()
 
     def tile(window: Window) -> Optional[Path]:
+        file_name = f"{metadata.stem}_{window.name()}.tif"
+        if file_name in existing_tiles:
+            return None
         rasterio_window = window.rasterio_window()
         with read_lock:
             data = dataset.read(1, window=rasterio_window)
@@ -97,7 +113,7 @@ def tile(window: Window) -> Optional[Path]:
             "transform": transform,
             "crs": dataset.crs,
         }
-        path = directory / f"{metadata.stem}_{window.name()}.tif"
+        path = directory / file_name
         with MemoryFile() as memory_file:
             with memory_file.open(**profile) as open_memory_file:
                 open_memory_file.write(data, 1)

diff --git a/tests/test_tile.py b/tests/test_tile.py
@@ -8,6 +8,13 @@ def test_tile_cdl(cdl: Path, tmp_path: Path) -> None:
     assert len(paths) == 4
 
 
+def test_tile_cdl_existing_tiles(cdl: Path, tmp_path: Path) -> None:
+    paths = tile.tile_geotiff(
+        cdl, tmp_path, 500, existing_tiles=["2021_30m_cdls_-91095_1807575_15000.tif"]
+    )
+    assert len(paths) == 3
+
+
 def test_tile_confidence(confidence: Path, tmp_path: Path) -> None:
     paths = tile.tile_geotiff(confidence, tmp_path, 500)
     assert len(paths) == 4