Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 15 additions & 4 deletions xrspatial/geotiff/_writers/eager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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``
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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'):
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 11 additions & 1 deletion xrspatial/geotiff/_writers/gpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -512,5 +521,6 @@ def _gpu_compress_to_part(gpu_arr, w, h, spp):
)

_write_bytes(file_bytes, path)
return path


182 changes: 182 additions & 0 deletions xrspatial/geotiff/tests/test_writer_return_path_1938.py
Original file line number Diff line number Diff line change
@@ -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
Loading