Add pyogrio stubs (#16004)

Author: hamdanal

lib/ts_utils/py315.py

@@ -8,6 +8,7 @@
     "hnswlib",
     "networkx",
     "pycocotools",
+    "pyogrio",
     "resampy",
     "shapely",
     "tensorflow",

pyrightconfig.stricter.json

@@ -81,6 +81,7 @@
         "stubs/pyasn1",
         "stubs/Pygments",
         "stubs/PyMySQL",
+        "stubs/pyogrio",
         "stubs/python-jose",
         "stubs/pywin32",
         "stubs/PyYAML",

stubs/pyogrio/@tests/stubtest_allowlist.txt

@@ -0,0 +1,2 @@
+pyogrio\.tests(\..*)?
+pyogrio\._typing  # stubs only module

stubs/pyogrio/METADATA.toml

@@ -0,0 +1,4 @@
+version = "0.13.*"
+# Requires a version of numpy with a `py.typed` file
+dependencies = ["numpy>=1.20", "types-geopandas"]
+upstream-repository = "https://github.com/geopandas/pyogrio"

stubs/pyogrio/pyogrio/__init__.pyi

@@ -0,0 +1,47 @@
+from .core import (
+    __gdal_geos_version__ as __gdal_geos_version__,
+    __gdal_version__ as __gdal_version__,
+    __gdal_version_string__ as __gdal_version_string__,
+    detect_write_driver as detect_write_driver,
+    get_gdal_config_option as get_gdal_config_option,
+    get_gdal_data_path as get_gdal_data_path,
+    list_drivers as list_drivers,
+    list_drivers_details as list_drivers_details,
+    list_layers as list_layers,
+    read_bounds as read_bounds,
+    read_info as read_info,
+    set_gdal_config_options as set_gdal_config_options,
+    vsi_curl_clear_cache as vsi_curl_clear_cache,
+    vsi_listtree as vsi_listtree,
+    vsi_rmtree as vsi_rmtree,
+    vsi_unlink as vsi_unlink,
+)
+from .geopandas import read_dataframe as read_dataframe, write_dataframe as write_dataframe
+from .raw import open_arrow as open_arrow, read_arrow as read_arrow, write_arrow as write_arrow
+
+__all__ = [
+    "__gdal_geos_version__",
+    "__gdal_version__",
+    "__gdal_version_string__",
+    "__version__",
+    "detect_write_driver",
+    "get_gdal_config_option",
+    "get_gdal_data_path",
+    "list_drivers",
+    "list_drivers_details",
+    "list_layers",
+    "open_arrow",
+    "read_arrow",
+    "read_bounds",
+    "read_dataframe",
+    "read_info",
+    "set_gdal_config_options",
+    "vsi_curl_clear_cache",
+    "vsi_listtree",
+    "vsi_rmtree",
+    "vsi_unlink",
+    "write_arrow",
+    "write_dataframe",
+]
+
+__version__: str

stubs/pyogrio/pyogrio/_typing.pyi

@@ -0,0 +1,28 @@
+import io
+from _typeshed import SupportsRead
+from collections.abc import Collection
+from pathlib import Path
+from typing import Any, Protocol, TypeAlias, TypeVar, type_check_only
+from typing_extensions import CapsuleType
+
+import numpy as np
+
+_T = TypeVar("_T")
+_G = TypeVar("_G", bound=np.generic)
+_G_co = TypeVar("_G_co", bound=np.generic, covariant=True)
+
+@type_check_only
+class SupportsArrowCStream(Protocol):
+    def __arrow_c_stream__(self, requested_schema: object | None = None) -> CapsuleType: ...
+
+@type_check_only
+class SupportsArray(Protocol[_G_co]):
+    def __array__(self) -> np.ndarray[Any, np.dtype[_G_co]]: ...
+
+Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]]
+Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_G]]
+ReadPathOrBuffer: TypeAlias = str | Path | bytes | SupportsRead[bytes]
+WritePathOrBuffer: TypeAlias = str | Path | io.BytesIO
+
+DualArrayLike: TypeAlias = SupportsArray[_G] | Collection[_T] | _T
+ArrayLikeInt: TypeAlias = DualArrayLike[np.bool | np.integer, int]

stubs/pyogrio/pyogrio/core.pyi

@@ -0,0 +1,82 @@
+from pathlib import Path
+from typing import Any, Literal, TypedDict, type_check_only
+
+import numpy as np
+import shapely as shp
+
+from ._typing import Array1D, Array2D, ReadPathOrBuffer
+
+__gdal_version__: tuple[int, int, int]
+__gdal_version_string__: str
+__gdal_geos_version__: tuple[int, int, int] | None
+
+@type_check_only
+class _Capabilities(TypedDict):
+    random_read: bool
+    fast_set_next_by_index: bool
+    fast_spatial_filter: bool
+    fast_feature_count: bool
+    fast_total_bounds: bool
+
+@type_check_only
+class _LayerInfo(TypedDict):
+    layer_name: str
+    # crs is `None` for non-spatial layers
+    crs: str | None
+    fields: Array1D[np.object_]  # field names (strings)
+    dtypes: Array1D[np.object_]  # field dtypes (strings)
+    ogr_types: list[str]
+    ogr_subtypes: list[str]
+    encoding: str
+    fid_column: str
+    geometry_name: str
+    # geometry_type is `None` for non-spatial layers
+    geometry_type: str | None
+    features: int
+    # total_bounds is `None` for non-spatial layers or if expensive to compute
+    total_bounds: tuple[float, float, float, float] | None
+    driver: str
+    capabilities: _Capabilities
+    dataset_metadata: dict[str, str] | None
+    layer_metadata: dict[str, str] | None
+
+@type_check_only
+class _DriverDetails(TypedDict):
+    long_name: str
+    read: bool
+    append: bool
+    write: bool
+    supports_vsi: bool
+    help_topic_url: str | None
+    extensions: list[str] | None
+
+def list_drivers(read: bool = False, write: bool = False, append: bool = False) -> dict[str, Literal["r", "rw"]]: ...
+def list_drivers_details() -> dict[str, _DriverDetails]: ...
+def detect_write_driver(path: str | Path) -> str: ...  # `path` is coerced to string internally
+def list_layers(path_or_buffer: ReadPathOrBuffer, /) -> Array2D[np.object_]: ...
+def read_bounds(
+    path_or_buffer: ReadPathOrBuffer,
+    /,
+    layer: int | str | None = None,
+    skip_features: int = 0,
+    max_features: int | None = None,
+    where: str | None = None,
+    bbox: tuple[float, float, float, float] | None = None,
+    mask: shp.Geometry | None = None,
+) -> tuple[Array1D[np.int64], Array2D[np.float64]]: ...
+def read_info(
+    path_or_buffer: ReadPathOrBuffer,
+    /,
+    layer: int | str | None = None,
+    encoding: str | None = None,
+    force_feature_count: bool = False,
+    force_total_bounds: bool = False,
+    **kwargs: Any,  # Dataset open options passed to OGR
+) -> _LayerInfo: ...
+def set_gdal_config_options(options: dict[str, Any]) -> None: ...
+def get_gdal_config_option(name: str) -> Any: ...  # Could return str, int, bool, or None
+def get_gdal_data_path() -> str: ...
+def vsi_listtree(path: str | Path, pattern: str | None = None) -> list[str]: ...
+def vsi_rmtree(path: str | Path) -> None: ...
+def vsi_unlink(path: str | Path) -> None: ...
+def vsi_curl_clear_cache(prefix: str | Path = "") -> None: ...

stubs/pyogrio/pyogrio/errors.pyi

@@ -0,0 +1,6 @@
+class DataSourceError(RuntimeError): ...
+class DataLayerError(RuntimeError): ...
+class CRSError(DataLayerError): ...
+class FeatureError(DataLayerError): ...
+class GeometryError(DataLayerError): ...
+class FieldError(DataLayerError): ...

stubs/pyogrio/pyogrio/geopandas.pyi

@@ -0,0 +1,80 @@
+import os
+from collections.abc import Collection, Mapping
+from typing import Any, Literal, overload
+
+import geopandas as gpd
+import pandas as pd
+import shapely as shp
+
+from ._typing import ArrayLikeInt, ReadPathOrBuffer, WritePathOrBuffer
+
+@overload
+def read_dataframe(
+    path_or_buffer: ReadPathOrBuffer | os.PathLike[str],
+    /,
+    layer: int | str | None = None,
+    encoding: str | None = None,
+    columns: Collection[str] | None = None,
+    read_geometry: Literal[True] = True,
+    force_2d: bool = False,
+    skip_features: int = 0,
+    max_features: int | None = None,
+    where: str | None = None,
+    bbox: tuple[float, float, float, float] | None = None,
+    mask: shp.Geometry | None = None,
+    fids: ArrayLikeInt | None = None,
+    sql: str | None = None,
+    sql_dialect: str | None = None,
+    fid_as_index: bool = False,
+    use_arrow: bool | None = None,
+    on_invalid: Literal["raise", "warn", "ignore", "fix"] = "raise",
+    arrow_to_pandas_kwargs: Mapping[str, Any] | None = None,
+    datetime_as_string: bool = False,
+    mixed_offsets_as_utc: bool = True,
+    **kwargs: Any,  # Dataset open options passed to OGR
+) -> gpd.GeoDataFrame: ...
+@overload
+def read_dataframe(
+    path_or_buffer: ReadPathOrBuffer | os.PathLike[str],
+    /,
+    layer: int | str | None = None,
+    encoding: str | None = None,
+    columns: Collection[str] | None = None,
+    *,
+    read_geometry: Literal[False],
+    force_2d: bool = False,
+    skip_features: int = 0,
+    max_features: int | None = None,
+    where: str | None = None,
+    bbox: tuple[float, float, float, float] | None = None,
+    mask: shp.Geometry | None = None,
+    fids: ArrayLikeInt | None = None,
+    sql: str | None = None,
+    sql_dialect: str | None = None,
+    fid_as_index: bool = False,
+    use_arrow: bool | None = None,
+    on_invalid: Literal["raise", "warn", "ignore", "fix"] = "raise",
+    arrow_to_pandas_kwargs: Mapping[str, Any] | None = None,
+    datetime_as_string: bool = False,
+    mixed_offsets_as_utc: bool = True,
+    **kwargs: Any,  # Dataset open options passed to OGR
+) -> pd.DataFrame: ...
+
+def write_dataframe(
+    df: pd.DataFrame,
+    path: WritePathOrBuffer,
+    layer: str | None = None,
+    driver: str | None = None,
+    encoding: str | None = None,
+    geometry_type: str | None = None,
+    promote_to_multi: bool | None = None,
+    nan_as_null: bool = True,
+    append: bool = False,
+    use_arrow: bool | None = None,
+    dataset_metadata: dict[str, Any] | None = None,
+    layer_metadata: dict[str, Any] | None = None,
+    metadata: dict[str, Any] | None = None,
+    dataset_options: dict[str, Any] | None = None,
+    layer_options: dict[str, Any] | None = None,
+    **kwargs: Any,  # Additional driver-specific dataset or layer creation options passed to OGR
+) -> None: ...