diff --git a/xrspatial/geotiff/_writers/eager.py b/xrspatial/geotiff/_writers/eager.py index 650a6da7..90bc8f0c 100644 --- a/xrspatial/geotiff/_writers/eager.py +++ b/xrspatial/geotiff/_writers/eager.py @@ -65,7 +65,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray, max_z_error: float = 0.0, photometric: str | int = 'auto', allow_internal_only_jpeg: bool = False, - allow_unparseable_crs: bool = False) -> None: + allow_unparseable_crs: bool = False) -> str | BinaryIO: """Write data as a GeoTIFF or Cloud Optimized GeoTIFF. Dask-backed DataArrays are written in streaming mode: one tile-row @@ -217,6 +217,16 @@ def to_geotiff(data: xr.DataArray | np.ndarray, ``"EPSG:4326"`` token on a host without pyproj produces a citation that most readers cannot interpret. See issue #1929. + 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 @@ -390,7 +400,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray, max_z_error=max_z_error, photometric=photometric, allow_unparseable_crs=allow_unparseable_crs) - return + return path # Dispatch to write_geotiff_gpu when GPU was selected (explicit # ``gpu=True`` or auto-detected CuPy data). ``auto_detected_gpu`` @@ -436,7 +446,7 @@ def to_geotiff(data: xr.DataArray | np.ndarray, allow_internal_only_jpeg=allow_internal_only_jpeg, allow_unparseable_crs=allow_unparseable_crs, ) - 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: @@ -593,7 +603,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'): @@ -679,6 +689,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, diff --git a/xrspatial/geotiff/_writers/gpu.py b/xrspatial/geotiff/_writers/gpu.py index a42a8b97..9a356838 100644 --- a/xrspatial/geotiff/_writers/gpu.py +++ b/xrspatial/geotiff/_writers/gpu.py @@ -47,7 +47,8 @@ def write_geotiff_gpu(data: xr.DataArray | cupy.ndarray | np.ndarray, max_z_error: float = 0.0, photometric: str | int = 'auto', allow_internal_only_jpeg: bool = False, - allow_unparseable_crs: bool = False) -> None: + allow_unparseable_crs: 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 @@ -179,6 +180,14 @@ def write_geotiff_gpu(data: xr.DataArray | cupy.ndarray | np.ndarray, :func:`to_geotiff` for the full description; the GPU writer applies the same fail-closed default. See issue #1929. + 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 @@ -512,5 +521,6 @@ def _gpu_compress_to_part(gpu_arr, w, h, spp): ) _write_bytes(file_bytes, path) + return path diff --git a/xrspatial/geotiff/tests/test_writer_return_path_1938.py b/xrspatial/geotiff/tests/test_writer_return_path_1938.py new file mode 100644 index 00000000..f364f463 --- /dev/null +++ b/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