Skip to content

Commit

Permalink
Fix Unexpected shape error with errors_as_nodata (#256)
Browse files Browse the repository at this point in the history
Opted not to refactor `NodataReader` as described in comments because #232 removes it all together, and I don't need even more merge conflicts someday.

Fixes #243
  • Loading branch information
gjoseph92 authored Aug 10, 2024
1 parent 7ca5eac commit 75919e3
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 10 deletions.
3 changes: 3 additions & 0 deletions stackstac/nodata_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@
State = Tuple[np.dtype, Union[int, float]]


# NOTE: this really should be a `ThreadsafeRioDataset` in `rio_reader.py`,
# not a `Reader` (it's never used as one).
class NodataReader:
"Reader that returns a constant (nodata) value for all reads"

scale_offset = (1.0, 0.0)

def __init__(
Expand Down
24 changes: 14 additions & 10 deletions stackstac/rio_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ def _curthread():
open=dict(
GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR",
# ^ stop GDAL from requesting `.aux` and `.msk` files from the bucket (speeds up `open` time a lot)
VSI_CACHE=True
VSI_CACHE=True,
# ^ cache HTTP requests for opening datasets. This is critical for `ThreadLocalRioDataset`,
# which re-opens the same URL many times---having the request cached makes subsequent `open`s
# in different threads snappy.
Expand Down Expand Up @@ -70,11 +70,9 @@ def _curthread():
class ThreadsafeRioDataset(Protocol):
scale_offset: Tuple[Union[int, float], Union[int, float]]

def read(self, window: Window, **kwargs) -> np.ndarray:
...
def read(self, window: Window, **kwargs) -> np.ndarray: ...

def close(self) -> None:
...
def close(self) -> None: ...


class SingleThreadedRioDataset:
Expand Down Expand Up @@ -408,8 +406,14 @@ def read(self, window: Window, **kwargs) -> np.ndarray:
result = np.ma.masked_array(result[0], mask=result[1] == 0)
elif result.shape[0] == 1:
result = result[0]
else:
raise RuntimeError(f"Unexpected shape {result.shape}, expected exactly 1 band.")
elif result.ndim != 2:
# We should only be getting `result.ndim == 2` in the case when `_open` produced a `NodataReader`.
# `Reader`s always return 2D arrays, whereas `rasterio.read` returns 3D. Pedantically, `NodataReader`
# shouldn't be a `Reader`, but a `ThreadsafeRioDataset`, and it should return a 3D array,
# just to be more consistent.
raise RuntimeError(
f"Unexpected shape {result.shape}, expected exactly 1 band."
)

scale, offset = self.scale_offset

Expand All @@ -419,9 +423,9 @@ def read(self, window: Window, **kwargs) -> np.ndarray:
result += offset

result = np.ma.filled(result, fill_value=self.fill_value)
assert np.issubdtype(result.dtype, self.dtype), (
f"Expected result array with dtype {self.dtype!r}, got {result.dtype!r}"
)
assert np.issubdtype(
result.dtype, self.dtype
), f"Expected result array with dtype {self.dtype!r}, got {result.dtype!r}"
return result

def close(self) -> None:
Expand Down

0 comments on commit 75919e3

Please sign in to comment.