geotiff: writers return the written path (#1938)

to_geotiff and write_geotiff_gpu returned None while write_vrt returned
str. The drift broke mypy consumers who handled the three writers
uniformly and made the Sphinx-rendered docs surface inconsistent.

Make to_geotiff and write_geotiff_gpu return the path argument (a str
for filesystem paths, the file-like object for BytesIO destinations).
The change is additive: existing callers that discarded the previous
None return are unaffected, and chained writes (path = to_geotiff(da,
out); da2 = open_geotiff(path)) work without round-tripping through a
variable.

write_vrt's existing str return is preserved.

Regression test pins the runtime return value for to_geotiff and
write_vrt across str + BytesIO + COG + dask-streaming destinations,
plus a GPU branch gated on cupy+CUDA. A signature-level assertion pins
the return annotations of all three writers so a future change cannot
silently revert the contract.

Author: brendancol

xrspatial/geotiff/_writers/eager.py

@@ -64,7 +64,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray,
                streaming_buffer_bytes: int = 256 * 1024 * 1024,
                max_z_error: float = 0.0,
                photometric: str | int = 'auto',
-               allow_internal_only_jpeg: bool = False) -> None:
+               allow_internal_only_jpeg: bool = False) -> str | BinaryIO:
     """Write data as a GeoTIFF or Cloud Optimized GeoTIFF.
 
     Dask-backed DataArrays are written in streaming mode: one tile-row
@@ -204,6 +204,16 @@ def to_geotiff(data: xr.DataArray | np.ndarray,
         GPU dispatch path so callers can reach the same experimental
         encode via ``to_geotiff(..., gpu=True)``. See issue #1845.
 
+    Returns
+    -------
+    str or binary file-like
+        The ``path`` argument (a string for filesystem paths, the
+        file-like object for BytesIO destinations). Returning the path
+        lines up with ``write_vrt`` and lets callers chain a write into
+        a read without round-tripping through a variable; existing
+        callers that discarded the previous ``None`` return are
+        unaffected. See issue #1938.
+
     Raises
     ------
     ValueError
@@ -376,7 +386,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray,
                          bigtiff=bigtiff,
                          max_z_error=max_z_error,
                          photometric=photometric)
-        return
+        return path
 
     # Dispatch to write_geotiff_gpu when GPU was selected (explicit
     # ``gpu=True`` or auto-detected CuPy data). ``auto_detected_gpu``
@@ -421,7 +431,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray,
                 photometric=photometric,
                 allow_internal_only_jpeg=allow_internal_only_jpeg,
             )
-            return
+            return path
         except ImportError as e:
             # ``write_geotiff_gpu`` raises ImportError when cupy itself
             # can't be imported. nvCOMP absence doesn't surface here:
@@ -572,7 +582,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray,
                 max_z_error=max_z_error,
                 photometric=photometric,
             )
-            return
+            return path
 
         # Eager compute (numpy, CuPy, or dask+COG)
         if hasattr(raw, 'get'):
@@ -654,6 +664,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray,
         max_z_error=max_z_error,
         photometric=photometric,
     )
+    return path
 
 
 def _write_single_tile(chunk_data, path, geo_transform, epsg, wkt,

xrspatial/geotiff/_writers/gpu.py

@@ -46,7 +46,8 @@ def write_geotiff_gpu(data: xr.DataArray | cupy.ndarray | np.ndarray,
                       max_z_error: float = 0.0,
                       streaming_buffer_bytes: int = 256 * 1024 * 1024,
                       photometric: str | int = 'auto',
-                      allow_internal_only_jpeg: bool = False) -> None:
+                      allow_internal_only_jpeg: bool = False
+                      ) -> str | BinaryIO:
     """Write a CuPy-backed DataArray as a GeoTIFF with GPU compression.
 
     Tiles are extracted and compressed on the GPU via nvCOMP, then
@@ -173,6 +174,14 @@ def write_geotiff_gpu(data: xr.DataArray | cupy.ndarray | np.ndarray,
         the flag, ``compression='jpeg'`` raises ``ValueError`` for
         parity with ``to_geotiff``. See issue #1845.
 
+    Returns
+    -------
+    str or binary file-like
+        The ``path`` argument (a string for filesystem paths, the
+        file-like object for BytesIO destinations). Returning the path
+        mirrors ``to_geotiff`` and ``write_vrt`` so callers can handle
+        the three writers uniformly. See issue #1938.
+
     Raises
     ------
     ValueError
@@ -502,5 +511,6 @@ def _gpu_compress_to_part(gpu_arr, w, h, spp):
     )
 
     _write_bytes(file_bytes, path)
+    return path
 
 

xrspatial/geotiff/tests/test_writer_return_path_1938.py

@@ -0,0 +1,182 @@
+"""Regression test for #1938: writer entry points return the written path.
+
+``write_vrt`` returned ``str`` while ``to_geotiff`` and
+``write_geotiff_gpu`` returned ``None``. The drift broke ``mypy``
+consumers who handle the three writers uniformly and made the
+Sphinx-rendered docs surface inconsistent.
+
+This module asserts:
+
+1. ``to_geotiff`` returns the ``path`` argument for filesystem and
+   file-like destinations.
+2. ``write_geotiff_gpu``'s annotation matches the canonical ``path``
+   return (the runtime check is gated on cupy + CUDA availability and
+   skipped here so the CPU test suite stays green).
+3. ``write_vrt`` keeps returning the path (already conformant).
+4. The three entry points share the same ``Returns`` annotation in
+   ``inspect.signature``.
+"""
+from __future__ import annotations
+
+import importlib.util
+import inspect
+import io
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from xrspatial.geotiff import (
+    to_geotiff,
+    write_geotiff_gpu,
+    write_vrt,
+)
+
+
+def _gpu_available() -> bool:
+    if importlib.util.find_spec("cupy") is None:
+        return False
+    try:
+        import cupy
+
+        return bool(cupy.cuda.is_available())
+    except Exception:
+        return False
+
+
+_HAS_GPU = _gpu_available()
+_gpu_only = pytest.mark.skipif(
+    not _HAS_GPU, reason="cupy + CUDA required",
+)
+
+
+def _small_da() -> xr.DataArray:
+    arr = np.arange(16, dtype=np.float32).reshape(4, 4)
+    return xr.DataArray(
+        arr,
+        dims=("y", "x"),
+        coords={"y": np.arange(4)[::-1].astype(np.float64),
+                "x": np.arange(4).astype(np.float64)},
+        attrs={"crs": 4326},
+    )
+
+
+def test_to_geotiff_returns_string_path(tmp_path):
+    """``to_geotiff`` returns the str path passed in."""
+    da = _small_da()
+    out = tmp_path / "test_1938_str.tif"
+    rv = to_geotiff(da, str(out))
+    assert isinstance(rv, str), (
+        f"to_geotiff(str) must return a str, got {type(rv).__name__}"
+    )
+    assert rv == str(out)
+    assert os.path.exists(rv)
+
+
+def test_to_geotiff_returns_file_like(tmp_path):
+    """``to_geotiff`` returns the file-like object passed in."""
+    da = _small_da()
+    buf = io.BytesIO()
+    rv = to_geotiff(da, buf)
+    assert rv is buf, (
+        f"to_geotiff(BytesIO) must return the same file-like, "
+        f"got {type(rv).__name__}"
+    )
+    # The buffer was actually written to.
+    assert buf.tell() > 0 or len(buf.getvalue()) > 0
+
+
+def test_to_geotiff_cog_returns_path(tmp_path):
+    """COG path also returns the str path."""
+    da = _small_da()
+    out = tmp_path / "test_1938_cog.tif"
+    rv = to_geotiff(da, str(out), cog=True, tile_size=16)
+    assert isinstance(rv, str)
+    assert rv == str(out)
+    assert os.path.exists(rv)
+
+
+def test_to_geotiff_dask_streaming_returns_path(tmp_path):
+    """Dask-streaming write path also returns the str path."""
+    import dask.array as da_arr
+
+    arr = da_arr.arange(256, dtype=np.float32, chunks=64).reshape(16, 16)
+    da = xr.DataArray(
+        arr,
+        dims=("y", "x"),
+        coords={"y": np.arange(16)[::-1].astype(np.float64),
+                "x": np.arange(16).astype(np.float64)},
+        attrs={"crs": 4326},
+    )
+    out = tmp_path / "test_1938_dask.tif"
+    rv = to_geotiff(da, str(out))
+    assert isinstance(rv, str)
+    assert rv == str(out)
+    assert os.path.exists(rv)
+
+
+def test_write_vrt_returns_string_path(tmp_path):
+    """``write_vrt`` (already conformant) keeps returning the str path."""
+    # Create a source tif first.
+    src = tmp_path / "src.tif"
+    to_geotiff(_small_da(), str(src))
+    vrt_path = tmp_path / "out.vrt"
+    rv = write_vrt(str(vrt_path), [str(src)])
+    assert isinstance(rv, str)
+    assert rv == str(vrt_path)
+    assert os.path.exists(rv)
+
+
+@_gpu_only
+def test_write_geotiff_gpu_returns_string_path(tmp_path):
+    """GPU writer returns the str path (only runs with cupy + CUDA)."""
+    import cupy
+
+    arr_cpu = np.arange(16, dtype=np.float32).reshape(4, 4)
+    arr_gpu = cupy.asarray(arr_cpu)
+    da = xr.DataArray(
+        arr_gpu,
+        dims=("y", "x"),
+        coords={"y": np.arange(4)[::-1].astype(np.float64),
+                "x": np.arange(4).astype(np.float64)},
+        attrs={"crs": 4326},
+    )
+    out = tmp_path / "test_1938_gpu.tif"
+    rv = write_geotiff_gpu(da, str(out))
+    assert isinstance(rv, str)
+    assert rv == str(out)
+    assert os.path.exists(rv)
+
+
+def test_writer_signatures_declare_path_return():
+    """All three writers annotate the same return type.
+
+    The annotation is a string under ``from __future__ import annotations``;
+    pin the literal so the three writers cannot drift apart silently.
+    """
+    expected = {
+        to_geotiff: "str | BinaryIO",
+        write_geotiff_gpu: "str | BinaryIO",
+        write_vrt: "str",
+    }
+    for fn, expected_ann in expected.items():
+        sig = inspect.signature(fn)
+        assert sig.return_annotation == expected_ann, (
+            f"{fn.__name__} return annotation drifted: expected "
+            f"{expected_ann!r}, got {sig.return_annotation!r}"
+        )
+
+
+def test_writer_returns_are_not_none():
+    """None of the public writers may go back to returning ``None``."""
+    da = _small_da()
+    with tempfile.TemporaryDirectory() as td:
+        out = os.path.join(td, "out.tif")
+        rv = to_geotiff(da, out)
+        assert rv is not None
+        src = os.path.join(td, "src.tif")
+        to_geotiff(da, src)
+        vrt_rv = write_vrt(os.path.join(td, "m.vrt"), [src])
+        assert vrt_rv is not None