Output COG metadata when creating empty COG

odc/geo/_cog.py

@@ -8,6 +8,7 @@
 import itertools
 import warnings
 from contextlib import contextmanager
+from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
 from typing import Any, Dict, Generator, Iterable, List, Literal, Optional, Tuple, Union
@@ -25,11 +26,28 @@
 from .types import MaybeNodata, Shape2d, SomeShape, Unset, shape_, wh_
 from .warp import resampling_s2rio
 
-# pylint: disable=too-many-locals,too-many-branches,too-many-arguments,too-many-statements
+# pylint: disable=too-many-locals,too-many-branches,too-many-arguments,too-many-statements,too-many-instance-attributes
 
 AxisOrder = Union[Literal["YX"], Literal["YXS"], Literal["SYX"]]
 
 
+@dataclass
+class CogMeta:
+    """
+    COG metadata.
+    """
+
+    axis: AxisOrder
+    shape: Shape2d
+    tile: Shape2d
+    pix_shape: Tuple[int, ...]
+    dtype: Any
+    compression: int
+    predictor: Any
+    gbox: Optional[GeoBox] = None
+    overviews: Optional[List["CogMeta"]] = None
+
+
 def _without(cfg: Dict[str, Any], *skip: str) -> Dict[str, Any]:
     skip = set(skip)
     return {k: v for k, v in cfg.items() if k not in skip}
@@ -556,7 +574,7 @@ def make_empty_cog(
     predictor: Union[int, str, bool, Unset] = Unset(),
     blocksize: Union[int, List[Union[int, Tuple[int, int]]]] = 2048,
     **kw,
-) -> memoryview:
+) -> Tuple[CogMeta, memoryview]:
     # pylint: disable=import-outside-toplevel
     have.check_or_error("tifffile", "rasterio", "xarray")
     from tifffile import (
@@ -572,10 +590,10 @@ def make_empty_cog(
         compression = str(kw.pop("compress", "ADOBE_DEFLATE"))
     compression = compression.upper()
     compression = {"DEFLATE": "ADOBE_DEFLATE"}.get(compression, compression)
-    compression = COMPRESSION[compression]
+    _compression = int(COMPRESSION[compression])
 
     if isinstance(predictor, Unset):
-        predictor = compression not in TIFF.IMAGE_COMPRESSIONS
+        predictor = _compression not in TIFF.IMAGE_COMPRESSIONS
 
     if isinstance(blocksize, int):
         blocksize = [blocksize]
@@ -594,20 +612,14 @@ def make_empty_cog(
     else:
         nsamples = shape[0]
 
-    extratags: List[Tuple[int, int, int, Any]] = []
-    if gbox is not None:
-        extratags, _ = geotiff_metadata(
-            gbox, nodata=nodata, gdal_metadata=gdal_metadata
-        )
-
     buf = BytesIO()
 
     opts_common = {
         "dtype": dtype,
         "photometric": photometric,
         "planarconfig": planarconfig,
         "predictor": predictor,
-        "compression": compression,
+        "compression": _compression,
         "software": False,
         **kw,
     }
@@ -622,25 +634,45 @@ def _sh(shape: Shape2d) -> Tuple[int, ...]:
     tsz = _norm_blocksize(blocksize[-1])
     im_shape, _, nlevels = _compute_cog_spec(im_shape, tsz)
 
+    extratags: List[Tuple[int, int, int, Any]] = []
+    if gbox is not None:
+        gbox = gbox.expand(im_shape)
+        extratags, _ = geotiff_metadata(
+            gbox, nodata=nodata, gdal_metadata=gdal_metadata
+        )
+    # TODO: support nodata/gdal_metadata without gbox?
+
     _blocks = itertools.chain(iter(blocksize), itertools.repeat(blocksize[-1]))
 
     tw = TiffWriter(buf, bigtiff=True, shaped=False)
+    metas: List[CogMeta] = []
 
     for tsz, idx in zip(_blocks, range(nlevels + 1)):
+        pix_shape = _sh(im_shape)
+        tile = _norm_blocksize(tsz)
+        meta = CogMeta(
+            ax, im_shape, shape_(tile), pix_shape, dtype, _compression, predictor
+        )
+
         if idx == 0:
             kw = {**opts_common, "extratags": extratags}
+            meta.gbox = gbox
         else:
             kw = {**opts_common, "subfiletype": FILETYPE.REDUCEDIMAGE}
 
         tw.write(
             itertools.repeat(b""),
-            shape=_sh(im_shape),
-            tile=_norm_blocksize(tsz),
+            shape=pix_shape,
+            tile=tile,
             **kw,
         )
 
+        metas.append(meta)
         im_shape = im_shape.shrink2()
 
+    meta = metas[0]
+    meta.overviews = metas[1:]
+
     tw.close()
 
-    return buf.getbuffer()
+    return meta, buf.getbuffer()

tests/test_cog.py

@@ -150,14 +150,20 @@ def test_empty_cog(shape, blocksize, expect_ax, dtype, compression, expect_predi
         gbox = gbox.zoom_to(shape[:2])
         assert gbox.shape == shape[:2]
 
-    mm = make_empty_cog(
+    meta, mm = make_empty_cog(
         shape,
         dtype,
         gbox=gbox,
         blocksize=blocksize,
         compression=compression,
     )
     assert isinstance(mm, memoryview)
+    assert meta.axis == expect_ax
+    assert meta.dtype == dtype
+    assert meta.shape[0] >= gbox.shape[0]
+    assert meta.shape[1] >= gbox.shape[1]
+    assert meta.gbox is not None
+    assert meta.shape == meta.gbox.shape
 
     f = tifffile.TiffFile(BytesIO(mm))
     assert f.tiff.is_bigtiff