[rasterio] Accept WarpedVRT and __geo_interface__ geometries

stubs/rasterio/@tests/stubtest_allowlist.txt

@@ -3,7 +3,6 @@ rasterio\._typing
 rasterio\._affine_types
 
 # Stubs-only type aliases referenced in signatures only.
-rasterio\.features\.Geometry
 rasterio\.merge\.MethodFunction
 
 # Cython implementation-detail attributes auto-generated on every

stubs/rasterio/rasterio/_typing.pyi

@@ -1,12 +1,26 @@
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from enum import Enum
-from typing import Any, BinaryIO, TypeAlias
+from typing import Any, BinaryIO, Protocol, TypeAlias, type_check_only
 
 from rasterio.crs import CRS
-from rasterio.io import BufferedDatasetWriter, DatasetReader, DatasetWriter, MemoryFile
+from rasterio.io import DatasetReaderBase, MemoryFile
 from rasterio.windows import Window
 
-AnyDataset: TypeAlias = DatasetReader | DatasetWriter | BufferedDatasetWriter | MemoryFile
+# `DatasetReaderBase` covers every readable dataset handle: DatasetReader,
+# DatasetWriter, BufferedDatasetWriter, MemoryDataset, and WarpedVRT (via
+# WarpedVRTReaderBase). `MemoryFile` is a file wrapper, not a dataset.
+AnyDataset: TypeAlias = DatasetReaderBase | MemoryFile
+
+@type_check_only
+class _SupportsGeoInterface(Protocol):
+    @property
+    def __geo_interface__(self) -> Mapping[str, Any]: ...
+
+# A GeoJSON-like mapping, or any object exposing one through the
+# `__geo_interface__` protocol (e.g. shapely / geopandas geometries).
+# The runtime unwraps `__geo_interface__` before use, so both forms are
+# accepted anywhere a geometry is expected.
+Geometry: TypeAlias = Mapping[str, Any] | _SupportsGeoInterface  # noqa: Y047
 Colormap: TypeAlias = dict[int, tuple[int, int, int] | tuple[int, int, int, int]]
 CRSInput: TypeAlias = str | dict[str, str] | CRS
 FileOrBytes: TypeAlias = BinaryIO | bytes

stubs/rasterio/rasterio/_warp.pyi

@@ -1,10 +1,10 @@
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 from typing import Any, Final
 
 from numpy.typing import DTypeLike, NDArray
 from rasterio._affine_types import Affine
 from rasterio._io import DatasetReaderBase
-from rasterio._typing import CRSInput, Indexes, ShapeND, WindowInput, _GDALOption, _NestedScalar
+from rasterio._typing import CRSInput, Geometry, Indexes, ShapeND, WindowInput, _GDALOption, _NestedScalar
 from rasterio.control import GroundControlPoint
 from rasterio.crs import CRS
 from rasterio.enums import Resampling
@@ -16,7 +16,7 @@ DEFAULT_NODATA_FLAG: Final[object]
 
 def recursive_round(val: _NestedScalar, precision: int) -> _NestedScalar: ...
 def _transform_geom(
-    src_crs: CRSInput, dst_crs: CRSInput, geom: Mapping[str, Any] | Sequence[Mapping[str, Any]], precision: int
+    src_crs: CRSInput, dst_crs: CRSInput, geom: Geometry | Sequence[Geometry], precision: int
 ) -> dict[str, Any] | list[dict[str, Any]]: ...
 def _reproject(
     source: NDArray[Any] | Any,

stubs/rasterio/rasterio/features.pyi

@@ -1,20 +1,19 @@
 import logging
 import os
-from collections.abc import Iterable, Iterator, Mapping
-from typing import Any, Final, TypeAlias, overload
+from collections.abc import Iterable, Iterator
+from typing import Any, Final, overload
 from typing_extensions import deprecated
 
 import numpy as np
 from numpy.typing import DTypeLike, NDArray
 from rasterio._affine_types import Affine
+from rasterio._typing import Geometry as Geometry
 from rasterio.enums import MergeAlg as MergeAlg
-from rasterio.io import DatasetReader
+from rasterio.io import DatasetReaderBase
 from rasterio.windows import Window as Window
 
 log: Final[logging.Logger]
 
-Geometry: TypeAlias = Mapping[str, Any]
-
 def geometry_mask(
     geometries: Iterable[Geometry], out_shape: tuple[int, int], transform: Affine, all_touched: bool = False, invert: bool = False
 ) -> NDArray[np.bool_]: ...
@@ -44,15 +43,15 @@ def bounds(geometry: Geometry, north_up: bool = True, transform: Affine | None =
 
 @overload
 def geometry_window(
-    dataset: DatasetReader, shapes: Iterable[Geometry], pad_x: float = 0, pad_y: float = 0, *, boundless: bool = False
+    dataset: DatasetReaderBase, shapes: Iterable[Geometry], pad_x: float = 0, pad_y: float = 0, *, boundless: bool = False
 ) -> Window: ...
 @overload
 @deprecated(
     "`north_up`, `rotated`, and `pixel_precision` on features.geometry_window are "
     "unused since rasterio 1.2.1 and will be removed in a future release."
 )
 def geometry_window(
-    dataset: DatasetReader,
+    dataset: DatasetReaderBase,
     shapes: Iterable[Geometry],
     pad_x: float = 0,
     pad_y: float = 0,
@@ -64,7 +63,7 @@ def geometry_window(
 
 def is_valid_geom(geom: Geometry) -> bool: ...
 def dataset_features(
-    src: DatasetReader,
+    src: DatasetReaderBase,
     bidx: int | None = None,
     sampling: int = 1,
     band: bool = True,

stubs/rasterio/rasterio/mask.pyi

@@ -1,27 +1,28 @@
 import logging
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable
 from typing import Any, Final
 
 from numpy.typing import NDArray
 from rasterio._affine_types import Affine
+from rasterio._typing import Geometry
 from rasterio.errors import WindowError as WindowError
 from rasterio.features import geometry_mask as geometry_mask, geometry_window as geometry_window
-from rasterio.io import DatasetReader
+from rasterio.io import DatasetReaderBase
 
 logger: Final[logging.Logger]
 
 def raster_geometry_mask(
-    dataset: DatasetReader,
-    shapes: Iterable[Mapping[str, Any]],
+    dataset: DatasetReaderBase,
+    shapes: Iterable[Geometry],
     all_touched: bool = False,
     invert: bool = False,
     crop: bool = False,
     pad: bool = False,
     pad_width: float = 0.5,
 ) -> tuple[NDArray[Any], Affine, tuple[int, int, int, int]]: ...
 def mask(
-    dataset: DatasetReader,
-    shapes: Iterable[Mapping[str, Any]],
+    dataset: DatasetReaderBase,
+    shapes: Iterable[Geometry],
     all_touched: bool = False,
     invert: bool = False,
     nodata: float | None = None,

stubs/rasterio/rasterio/merge.pyi

@@ -7,7 +7,7 @@ from typing_extensions import deprecated
 from numpy.typing import DTypeLike, NDArray
 from rasterio._affine_types import Affine
 from rasterio.enums import Resampling
-from rasterio.io import DatasetReader
+from rasterio.io import DatasetReaderBase
 
 logger: Final[logging.Logger]
 
@@ -26,7 +26,7 @@ def copy_count(merged_data: _Arr, new_data: _Arr, merged_mask: _Arr, new_mask: _
 
 @overload
 def merge(
-    sources: Sequence[DatasetReader | str | os.PathLike[str]],
+    sources: Sequence[DatasetReaderBase | str | os.PathLike[str]],
     bounds: tuple[float, float, float, float] | None = None,
     res: float | tuple[float, float] | None = None,
     nodata: float | None = None,
@@ -46,7 +46,7 @@ def merge(
 @overload
 @deprecated("The `precision` parameter is unused since rasterio 1.3 and will be removed in 2.0.0.")
 def merge(
-    sources: Sequence[DatasetReader | str | os.PathLike[str]],
+    sources: Sequence[DatasetReaderBase | str | os.PathLike[str]],
     bounds: tuple[float, float, float, float] | None = None,
     res: float | tuple[float, float] | None = None,
     nodata: float | None = None,

stubs/rasterio/rasterio/sample.pyi

@@ -2,9 +2,12 @@ from collections.abc import Iterable, Iterator, Sequence
 from typing import Any
 
 from numpy.typing import NDArray
-from rasterio.io import DatasetReader
+from rasterio.io import DatasetReaderBase
 
 def sample_gen(
-    dataset: DatasetReader, xy: Iterable[tuple[float, float]], indexes: int | Sequence[int] | None = None, masked: bool = False
+    dataset: DatasetReaderBase,
+    xy: Iterable[tuple[float, float]],
+    indexes: int | Sequence[int] | None = None,
+    masked: bool = False,
 ) -> Iterator[NDArray[Any]]: ...
 def sort_xy(xy: Iterable[tuple[float, float]]) -> list[tuple[float, float]]: ...

stubs/rasterio/rasterio/stack.pyi

@@ -6,12 +6,12 @@ from typing import Any, Final
 from numpy.typing import DTypeLike, NDArray
 from rasterio._affine_types import Affine
 from rasterio.enums import Resampling
-from rasterio.io import DatasetReader
+from rasterio.io import DatasetReaderBase
 
 logger: Final[logging.Logger]
 
 def stack(
-    sources: Sequence[DatasetReader | str | os.PathLike[str]],
+    sources: Sequence[DatasetReaderBase | str | os.PathLike[str]],
     bounds: tuple[float, float, float, float] | None = None,
     res: float | tuple[float, float] | None = None,
     nodata: float | None = None,

stubs/rasterio/rasterio/warp.pyi

@@ -5,7 +5,7 @@ from typing_extensions import deprecated
 
 from numpy.typing import ArrayLike, NDArray
 from rasterio._affine_types import Affine
-from rasterio._typing import CRSInput, _GDALOption
+from rasterio._typing import CRSInput, Geometry, _GDALOption
 from rasterio.control import GroundControlPoint
 from rasterio.enums import Resampling
 from rasterio.rpc import RPC
@@ -22,7 +22,7 @@ def transform(
 
 @overload
 def transform_geom(
-    src_crs: CRSInput, dst_crs: CRSInput, geom: Mapping[str, Any] | Sequence[Mapping[str, Any]], *, precision: float = -1
+    src_crs: CRSInput, dst_crs: CRSInput, geom: Geometry | Sequence[Geometry], *, precision: float = -1
 ) -> dict[str, Any] | list[dict[str, Any]]: ...
 @overload
 @deprecated(
@@ -33,7 +33,7 @@ def transform_geom(
 def transform_geom(
     src_crs: CRSInput,
     dst_crs: CRSInput,
-    geom: Mapping[str, Any] | Sequence[Mapping[str, Any]],
+    geom: Geometry | Sequence[Geometry],
     antimeridian_cutting: bool | None = None,
     antimeridian_offset: float | None = None,
     precision: float = -1,