Offering two improvements to the open() type stub:

1. The DatasetReader and DatasetWriter stubs should support context manager ("with _ as _") syntax, achieved through the __enter__ and __exit__ methods.

   

2. Union[DatasetReader, DatasetWriter] is not quite the correct return type because it restricts usage to only methods contained in both classes. Instead, I overloaded the stub to return either a DatasetReader or DatasetWriter conditionally based on the mode parameter. This also results in more accurate type hinting.

   

Awesome work on this @kylebarron!

After installing and checking out some sample code, I'm getting the error referenced.

I see an Env stub exists in the source, and I can avoid the error by replacing my call with rio.env.Env. I also tried adding the import to get the original statement recognized, but I can't seem to get that working. Any idea what I'm doing wrong there?

Hey @kylebarron :wave: rasterio.path is going to be deprecated in 1.3 (see https://github.com/rasterio/rasterio/pull/2424) and this will affect https://github.com/kylebarron/types-rasterio/blob/142a531c9dd1a35672f1bd811fdf05356b241f02/rasterio-stubs/drivers.pyi#L3.

Enables a generic typevar to propagate through the class. For example, define a reader to be of type np.float64, and then data read from it will also be of that type.

from __future__ import annotations
from rasterio._io import DatasetReaderBase
import numpy as np

out: DatasetReaderBase[np.float64] = DatasetReaderBase('path')
output_array = out.read()
reveal_type(output_array)
# note: Revealed type is "numpy.ndarray[Any, numpy.dtype[numpy.floating*[numpy.typing._64Bit]]]"