Skip to content

Commit 0ff2bb2

Browse files
thomas-maschlerthomas-maschler
authored
[rasterio] Accept WarpedVRT and __geo_interface__ geometries (#16092)
Widen dataset parameters from the concrete DatasetReader to the shared DatasetReaderBase so any readable dataset (including WarpedVRT) is accepted, matching the runtime duck-typed contract: - merge.merge and stack.stack sources - features.geometry_window and features.dataset_features - mask.mask and mask.raster_geometry_mask - sample.sample_gen - the AnyDataset alias (now DatasetReaderBase | MemoryFile) Introduce a shared Geometry alias in _typing (a GeoJSON-like Mapping or any object implementing __geo_interface__, e.g. shapely geometries) and use it wherever the runtime unwraps __geo_interface__: - features.rasterize, geometry_mask, bounds, is_valid_geom, geometry_window - mask.mask and mask.raster_geometry_mask shapes - warp.transform_geom and _warp._transform_geom `Geometry` now lives in the blanket-allowlisted `rasterio._typing` module and is only re-exported from `rasterio.features`.
1 parent 2827b05 commit 0ff2bb2

9 files changed

Lines changed: 48 additions & 32 deletions

File tree

‎stubs/rasterio/@tests/stubtest_allowlist.txt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ rasterio\._typing
33
rasterio\._affine_types
44

55
# Stubs-only type aliases referenced in signatures only.
6-
rasterio\.features\.Geometry
76
rasterio\.merge\.MethodFunction
87

98
# Cython implementation-detail attributes auto-generated on every

‎stubs/rasterio/rasterio/_typing.pyi‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,26 @@
1-
from collections.abc import Callable, Sequence
1+
from collections.abc import Callable, Mapping, Sequence
22
from enum import Enum
3-
from typing import Any, BinaryIO, TypeAlias
3+
from typing import Any, BinaryIO, Protocol, TypeAlias, type_check_only
44

55
from rasterio.crs import CRS
6-
from rasterio.io import BufferedDatasetWriter, DatasetReader, DatasetWriter, MemoryFile
6+
from rasterio.io import DatasetReaderBase, MemoryFile
77
from rasterio.windows import Window
88

9-
AnyDataset: TypeAlias = DatasetReader | DatasetWriter | BufferedDatasetWriter | MemoryFile
9+
# `DatasetReaderBase` covers every readable dataset handle: DatasetReader,
10+
# DatasetWriter, BufferedDatasetWriter, MemoryDataset, and WarpedVRT (via
11+
# WarpedVRTReaderBase). `MemoryFile` is a file wrapper, not a dataset.
12+
AnyDataset: TypeAlias = DatasetReaderBase | MemoryFile
13+
14+
@type_check_only
15+
class _SupportsGeoInterface(Protocol):
16+
@property
17+
def __geo_interface__(self) -> Mapping[str, Any]: ...
18+
19+
# A GeoJSON-like mapping, or any object exposing one through the
20+
# `__geo_interface__` protocol (e.g. shapely / geopandas geometries).
21+
# The runtime unwraps `__geo_interface__` before use, so both forms are
22+
# accepted anywhere a geometry is expected.
23+
Geometry: TypeAlias = Mapping[str, Any] | _SupportsGeoInterface # noqa: Y047
1024
Colormap: TypeAlias = dict[int, tuple[int, int, int] | tuple[int, int, int, int]]
1125
CRSInput: TypeAlias = str | dict[str, str] | CRS
1226
FileOrBytes: TypeAlias = BinaryIO | bytes

‎stubs/rasterio/rasterio/_warp.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
from collections.abc import Mapping, Sequence
1+
from collections.abc import Sequence
22
from typing import Any, Final
33

44
from numpy.typing import DTypeLike, NDArray
55
from rasterio._affine_types import Affine
66
from rasterio._io import DatasetReaderBase
7-
from rasterio._typing import CRSInput, Indexes, ShapeND, WindowInput, _GDALOption, _NestedScalar
7+
from rasterio._typing import CRSInput, Geometry, Indexes, ShapeND, WindowInput, _GDALOption, _NestedScalar
88
from rasterio.control import GroundControlPoint
99
from rasterio.crs import CRS
1010
from rasterio.enums import Resampling
@@ -16,7 +16,7 @@ DEFAULT_NODATA_FLAG: Final[object]
1616

1717
def recursive_round(val: _NestedScalar, precision: int) -> _NestedScalar: ...
1818
def _transform_geom(
19-
src_crs: CRSInput, dst_crs: CRSInput, geom: Mapping[str, Any] | Sequence[Mapping[str, Any]], precision: int
19+
src_crs: CRSInput, dst_crs: CRSInput, geom: Geometry | Sequence[Geometry], precision: int
2020
) -> dict[str, Any] | list[dict[str, Any]]: ...
2121
def _reproject(
2222
source: NDArray[Any] | Any,

‎stubs/rasterio/rasterio/features.pyi‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,19 @@
11
import logging
22
import os
3-
from collections.abc import Iterable, Iterator, Mapping
4-
from typing import Any, Final, TypeAlias, overload
3+
from collections.abc import Iterable, Iterator
4+
from typing import Any, Final, overload
55
from typing_extensions import deprecated
66

77
import numpy as np
88
from numpy.typing import DTypeLike, NDArray
99
from rasterio._affine_types import Affine
10+
from rasterio._typing import Geometry as Geometry
1011
from rasterio.enums import MergeAlg as MergeAlg
11-
from rasterio.io import DatasetReader
12+
from rasterio.io import DatasetReaderBase
1213
from rasterio.windows import Window as Window
1314

1415
log: Final[logging.Logger]
1516

16-
Geometry: TypeAlias = Mapping[str, Any]
17-
1817
def geometry_mask(
1918
geometries: Iterable[Geometry], out_shape: tuple[int, int], transform: Affine, all_touched: bool = False, invert: bool = False
2019
) -> NDArray[np.bool_]: ...
@@ -44,15 +43,15 @@ def bounds(geometry: Geometry, north_up: bool = True, transform: Affine | None =
4443

4544
@overload
4645
def geometry_window(
47-
dataset: DatasetReader, shapes: Iterable[Geometry], pad_x: float = 0, pad_y: float = 0, *, boundless: bool = False
46+
dataset: DatasetReaderBase, shapes: Iterable[Geometry], pad_x: float = 0, pad_y: float = 0, *, boundless: bool = False
4847
) -> Window: ...
4948
@overload
5049
@deprecated(
5150
"`north_up`, `rotated`, and `pixel_precision` on features.geometry_window are "
5251
"unused since rasterio 1.2.1 and will be removed in a future release."
5352
)
5453
def geometry_window(
55-
dataset: DatasetReader,
54+
dataset: DatasetReaderBase,
5655
shapes: Iterable[Geometry],
5756
pad_x: float = 0,
5857
pad_y: float = 0,
@@ -64,7 +63,7 @@ def geometry_window(
6463

6564
def is_valid_geom(geom: Geometry) -> bool: ...
6665
def dataset_features(
67-
src: DatasetReader,
66+
src: DatasetReaderBase,
6867
bidx: int | None = None,
6968
sampling: int = 1,
7069
band: bool = True,

‎stubs/rasterio/rasterio/mask.pyi‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,28 @@
11
import logging
2-
from collections.abc import Iterable, Mapping
2+
from collections.abc import Iterable
33
from typing import Any, Final
44

55
from numpy.typing import NDArray
66
from rasterio._affine_types import Affine
7+
from rasterio._typing import Geometry
78
from rasterio.errors import WindowError as WindowError
89
from rasterio.features import geometry_mask as geometry_mask, geometry_window as geometry_window
9-
from rasterio.io import DatasetReader
10+
from rasterio.io import DatasetReaderBase
1011

1112
logger: Final[logging.Logger]
1213

1314
def raster_geometry_mask(
14-
dataset: DatasetReader,
15-
shapes: Iterable[Mapping[str, Any]],
15+
dataset: DatasetReaderBase,
16+
shapes: Iterable[Geometry],
1617
all_touched: bool = False,
1718
invert: bool = False,
1819
crop: bool = False,
1920
pad: bool = False,
2021
pad_width: float = 0.5,
2122
) -> tuple[NDArray[Any], Affine, tuple[int, int, int, int]]: ...
2223
def mask(
23-
dataset: DatasetReader,
24-
shapes: Iterable[Mapping[str, Any]],
24+
dataset: DatasetReaderBase,
25+
shapes: Iterable[Geometry],
2526
all_touched: bool = False,
2627
invert: bool = False,
2728
nodata: float | None = None,

‎stubs/rasterio/rasterio/merge.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ from typing_extensions import deprecated
77
from numpy.typing import DTypeLike, NDArray
88
from rasterio._affine_types import Affine
99
from rasterio.enums import Resampling
10-
from rasterio.io import DatasetReader
10+
from rasterio.io import DatasetReaderBase
1111

1212
logger: Final[logging.Logger]
1313

@@ -26,7 +26,7 @@ def copy_count(merged_data: _Arr, new_data: _Arr, merged_mask: _Arr, new_mask: _
2626

2727
@overload
2828
def merge(
29-
sources: Sequence[DatasetReader | str | os.PathLike[str]],
29+
sources: Sequence[DatasetReaderBase | str | os.PathLike[str]],
3030
bounds: tuple[float, float, float, float] | None = None,
3131
res: float | tuple[float, float] | None = None,
3232
nodata: float | None = None,
@@ -46,7 +46,7 @@ def merge(
4646
@overload
4747
@deprecated("The `precision` parameter is unused since rasterio 1.3 and will be removed in 2.0.0.")
4848
def merge(
49-
sources: Sequence[DatasetReader | str | os.PathLike[str]],
49+
sources: Sequence[DatasetReaderBase | str | os.PathLike[str]],
5050
bounds: tuple[float, float, float, float] | None = None,
5151
res: float | tuple[float, float] | None = None,
5252
nodata: float | None = None,

‎stubs/rasterio/rasterio/sample.pyi‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ from collections.abc import Iterable, Iterator, Sequence
22
from typing import Any
33

44
from numpy.typing import NDArray
5-
from rasterio.io import DatasetReader
5+
from rasterio.io import DatasetReaderBase
66

77
def sample_gen(
8-
dataset: DatasetReader, xy: Iterable[tuple[float, float]], indexes: int | Sequence[int] | None = None, masked: bool = False
8+
dataset: DatasetReaderBase,
9+
xy: Iterable[tuple[float, float]],
10+
indexes: int | Sequence[int] | None = None,
11+
masked: bool = False,
912
) -> Iterator[NDArray[Any]]: ...
1013
def sort_xy(xy: Iterable[tuple[float, float]]) -> list[tuple[float, float]]: ...

‎stubs/rasterio/rasterio/stack.pyi‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@ from typing import Any, Final
66
from numpy.typing import DTypeLike, NDArray
77
from rasterio._affine_types import Affine
88
from rasterio.enums import Resampling
9-
from rasterio.io import DatasetReader
9+
from rasterio.io import DatasetReaderBase
1010

1111
logger: Final[logging.Logger]
1212

1313
def stack(
14-
sources: Sequence[DatasetReader | str | os.PathLike[str]],
14+
sources: Sequence[DatasetReaderBase | str | os.PathLike[str]],
1515
bounds: tuple[float, float, float, float] | None = None,
1616
res: float | tuple[float, float] | None = None,
1717
nodata: float | None = None,

‎stubs/rasterio/rasterio/warp.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ from typing_extensions import deprecated
55

66
from numpy.typing import ArrayLike, NDArray
77
from rasterio._affine_types import Affine
8-
from rasterio._typing import CRSInput, _GDALOption
8+
from rasterio._typing import CRSInput, Geometry, _GDALOption
99
from rasterio.control import GroundControlPoint
1010
from rasterio.enums import Resampling
1111
from rasterio.rpc import RPC
@@ -22,7 +22,7 @@ def transform(
2222

2323
@overload
2424
def transform_geom(
25-
src_crs: CRSInput, dst_crs: CRSInput, geom: Mapping[str, Any] | Sequence[Mapping[str, Any]], *, precision: float = -1
25+
src_crs: CRSInput, dst_crs: CRSInput, geom: Geometry | Sequence[Geometry], *, precision: float = -1
2626
) -> dict[str, Any] | list[dict[str, Any]]: ...
2727
@overload
2828
@deprecated(
@@ -33,7 +33,7 @@ def transform_geom(
3333
def transform_geom(
3434
src_crs: CRSInput,
3535
dst_crs: CRSInput,
36-
geom: Mapping[str, Any] | Sequence[Mapping[str, Any]],
36+
geom: Geometry | Sequence[Geometry],
3737
antimeridian_cutting: bool | None = None,
3838
antimeridian_offset: float | None = None,
3939
precision: float = -1,

0 commit comments

Comments
 (0)