rasterio.open typing updates

[alecglen]

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.

   

[kylebarron]

Does this work? Usually you need an overload for each discriminator, and then a fallback, but maybe it works a little differently in stub files.

@overload
def open(
    fp,
    mode: str = "w",
    ...
    ) -> DatasetWriter: ...

@overload
def open(
    fp,
    mode: str = "r",
    ...
    ) -> DatasetReader: ...

def open(
    fp,
    mode: str = ...,
	...
    ) -> Union[DatasetReader, DatasetWriter]: ...

For example, does rasterio support anything other than w and r? Does it support an append mode?

[alecglen]

Good question, I was just getting it going for my use case and then figured I'd share the change. I'll take a look and update it to be more comprehensive.

[kylebarron]

These don't exist in the rasterio implementation of DatasetReader: https://github.com/rasterio/rasterio/blob/38ad052cc2ab460d0b632bd2f3d7919e3d96c3c8/rasterio/io.py#L28-L34

Ideally they'd be put on DatasetBase:
https://github.com/rasterio/rasterio/blob/38ad052cc2ab460d0b632bd2f3d7919e3d96c3c8/rasterio/_base.pyx#L436-L441