Add geopandas stubs (#12990)

Author: hamdanal

pyrightconfig.json

@@ -9,6 +9,7 @@
         // Stubs that don't work in all Python versions
         "stubs/seaborn",
         "stubs/shapely",
+        "stubs/geopandas",
         // test cases use a custom config file
         "**/@tests/test_cases",
     ],

pyrightconfig.stricter.json

@@ -49,6 +49,7 @@
         "stubs/Flask-SocketIO",
         "stubs/fpdf2",
         "stubs/gdb",
+        "stubs/geopandas",
         "stubs/google-cloud-ndb",
         "stubs/hdbcli/hdbcli/dbapi.pyi",
         "stubs/html5lib",

stubs/geopandas/@tests/stubtest_allowlist.txt

@@ -0,0 +1,24 @@
+# Stub missing OK
+geopandas\.conftest
+geopandas\.(.*\.)?tests.*
+geopandas\.array\.(type_mapping|geometry_type_ids|geometry_type_values)
+geopandas\.io\.file\.(FIONA|PYOGRIO)_GE_.*
+geopandas\.io\.file\.(fiona|pyogrio)(_env|_import_error)?
+geopandas\.io\.util
+geopandas\.datasets\.*
+
+# Inconsistent OK
+geopandas\.(geodataframe\.)?GeoDataFrame\.plot
+
+# Failed to import OK (require extra dependencies)
+geopandas\.io\._geoarrow
+geopandas\._exports
+
+# Not present in stub OK (do not work at runtime, to be removed from geopandas)
+geopandas\.(geoseries\.)?GeoSeries\.append # https://github.com/geopandas/geopandas/pull/3460
+geopandas\.(geoseries\.)?GeoSeries\.select # https://github.com/geopandas/geopandas/pull/3394
+
+# Inconsistent (TODO)
+geopandas\.(geoseries\.)?GeoSeries\.apply
+geopandas\.(geoseries\.)?GeoSeries\.fillna
+geopandas\.(geoseries\.)?GeoSeries\.sort_index

stubs/geopandas/METADATA.toml

@@ -0,0 +1,4 @@
+version = "1.0.1"
+# Requires a version of numpy with a `py.typed` file
+requires = ["numpy>=1.20", "pandas-stubs", "types-shapely", "pyproj"]
+upstream_repository = "https://github.com/geopandas/geopandas"

stubs/geopandas/geopandas/__init__.pyi

@@ -0,0 +1,20 @@
+from typing import Final
+
+from ._config import options as options
+from ._exports import (
+    gpd as gpd,
+    list_layers as list_layers,
+    np as np,
+    pd as pd,
+    read_feather as read_feather,
+    read_file as read_file,
+    read_parquet as read_parquet,
+    read_postgis as read_postgis,
+)
+from .array import points_from_xy as points_from_xy
+from .geodataframe import GeoDataFrame as GeoDataFrame
+from .geoseries import GeoSeries as GeoSeries
+from .tools import clip as clip, overlay as overlay, sjoin as sjoin, sjoin_nearest as sjoin_nearest
+from .tools._show_versions import show_versions as show_versions
+
+__version__: Final[str]

stubs/geopandas/geopandas/_config.pyi

@@ -0,0 +1,21 @@
+from _typeshed import SupportsItems, Unused
+from collections.abc import Callable
+from typing import Any, NamedTuple
+
+class Option(NamedTuple):
+    key: str
+    default_value: Any  # Can be "any" type
+    doc: str
+    validator: Callable[[object], Unused]
+    callback: Callable[[str, object], Unused] | None
+
+class Options:
+    def __init__(self, options: SupportsItems[str, Option]) -> None: ...
+    # Accept and return arbitrary values
+    def __setattr__(self, key: str, value: Any) -> None: ...
+    def __getattr__(self, key: str) -> Any: ...
+
+display_precision: Option
+use_pygeos: Option
+io_engine: Option
+options: Options

stubs/geopandas/geopandas/_decorator.pyi

@@ -0,0 +1,21 @@
+from collections.abc import Callable
+from typing import TypeVar, overload
+from typing_extensions import TypeAlias
+
+_AnyCallable: TypeAlias = Callable[..., object]
+_Func = TypeVar("_Func", bound=_AnyCallable)
+
+# We (ab)use this decorator to also copy the signature of the source function (overload 1)
+# The advantages are:
+# - avoid copying all parameters and types manually while conserving type safety
+# - signature properly handeled in IDEs (at least with Pylance)
+# - docstring from the original function properly displayed (at least with Pylance)
+# Using the other overloads returns the signature of the decorated function instead
+@overload
+def doc(func: _Func, /, **params: object) -> Callable[[_AnyCallable], _Func]: ...
+@overload
+def doc(docstring: str, /, *docstrings: str | _AnyCallable, **params: object) -> Callable[[_Func], _Func]: ...
+@overload
+def doc(
+    docstring1: str | _AnyCallable, docstring2: str | _AnyCallable, /, *docstrings: str | _AnyCallable, **params: object
+) -> Callable[[_Func], _Func]: ...

stubs/geopandas/geopandas/_exports.pyi

@@ -0,0 +1,21 @@
+# Type checking-only module to export public symbols with a different name in __init__.pyi
+
+import geopandas as gpd
+import numpy as np
+import pandas as pd
+from geopandas.io.arrow import _read_feather as read_feather, _read_parquet as read_parquet
+from geopandas.io.file import _list_layers as list_layers, _read_file as read_file
+from geopandas.io.sql import _read_postgis as read_postgis
+
+__all__ = [
+    # IO functions
+    "read_file",
+    "read_feather",
+    "read_parquet",
+    "read_postgis",
+    "list_layers",
+    # Modules for interactive use
+    "np",
+    "pd",
+    "gpd",
+]

stubs/geopandas/geopandas/array.pyi

@@ -0,0 +1,244 @@
+import builtins
+from _typeshed import Incomplete, Unused
+from collections.abc import Callable, Sequence
+from typing import Any, ClassVar, Literal, NoReturn, SupportsIndex, TypeVar, overload
+from typing_extensions import Self, TypeAlias, deprecated
+
+import numpy as np
+import pandas as pd
+from numpy.typing import ArrayLike, DTypeLike, NDArray
+from pandas.api.extensions import ExtensionArray, ExtensionDtype
+from pyproj import CRS, Transformer
+from shapely import Geometry
+from shapely.geometry.base import BaseGeometry
+
+from .base import _AffinityOrigin, _ConvertibleToCRS
+from .sindex import SpatialIndex
+
+_ScalarType = TypeVar("_ScalarType", bound=np.generic)
+_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_ScalarType]]
+_Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_ScalarType]]
+_ArrayOrGeom: TypeAlias = GeometryArray | ArrayLike | Geometry
+
+TransformerFromCRS = Transformer.from_crs
+
+class GeometryDtype(ExtensionDtype):
+    type: ClassVar[type[BaseGeometry]]
+    name: ClassVar[str]
+    na_value: float
+    @classmethod
+    def construct_from_string(cls, string: str) -> Self: ...
+    @classmethod
+    def construct_array_type(cls) -> builtins.type[GeometryArray]: ...
+
+def isna(value: object) -> bool: ...
+def from_shapely(data, crs: _ConvertibleToCRS | None = None) -> GeometryArray: ...
+def to_shapely(geoms: GeometryArray) -> _Array1D[np.object_]: ...
+def from_wkb(
+    data, crs: _ConvertibleToCRS | None = None, on_invalid: Literal["raise", "warn", "ignore"] = "raise"
+) -> GeometryArray: ...
+@overload
+def to_wkb(geoms: GeometryArray, hex: Literal[False] = False, **kwargs) -> _Array1D[np.bytes_]: ...
+@overload
+def to_wkb(geoms: GeometryArray, hex: Literal[True], **kwargs) -> _Array1D[np.str_]: ...
+def from_wkt(
+    data, crs: _ConvertibleToCRS | None = None, on_invalid: Literal["raise", "warn", "ignore"] = "raise"
+) -> GeometryArray: ...
+def to_wkt(geoms: GeometryArray, **kwargs) -> _Array1D[np.str_]: ...
+def points_from_xy(
+    x: ArrayLike, y: ArrayLike, z: ArrayLike | None = None, crs: _ConvertibleToCRS | None = None
+) -> GeometryArray: ...
+
+class GeometryArray(ExtensionArray):
+    def __init__(self, data: GeometryArray | NDArray[np.object_], crs: _ConvertibleToCRS | None = None) -> None: ...
+    @property
+    def sindex(self) -> SpatialIndex: ...
+    @property
+    def has_sindex(self) -> bool: ...
+    @property
+    def crs(self) -> CRS | None: ...
+    @crs.setter
+    def crs(self, value: _ConvertibleToCRS | None) -> None: ...
+    def check_geographic_crs(self, stacklevel: int) -> None: ...
+    @property
+    def dtype(self) -> GeometryDtype: ...
+    def __len__(self) -> int: ...
+    # np.integer[Any] because precision is not important
+    @overload
+    def __getitem__(self, idx: int | np.integer[Any]) -> BaseGeometry: ...  # Always 1-D, doesn't accept tuple
+    @overload
+    def __getitem__(
+        self, idx: slice | Sequence[SupportsIndex] | NDArray[np.bool_] | NDArray[np.integer[Any]]
+    ) -> GeometryArray: ...
+    @overload
+    def __getitem__(
+        self, idx: int | np.integer[Any] | slice | Sequence[int] | NDArray[np.bool_] | NDArray[np.integer[Any]]
+    ) -> BaseGeometry | GeometryArray: ...
+    def __setitem__(
+        self, key, value: _ArrayOrGeom | pd.DataFrame | pd.Series[Any]  # Cannot use pd.Series[BaseGeometry]
+    ) -> None: ...
+    @property
+    def is_valid(self) -> _Array1D[np.bool_]: ...
+    def is_valid_reason(self) -> _Array1D[np.object_]: ...
+    @property
+    def is_empty(self) -> _Array1D[np.bool_]: ...
+    @property
+    def is_simple(self) -> _Array1D[np.bool_]: ...
+    @property
+    def is_ring(self) -> _Array1D[np.bool_]: ...
+    @property
+    def is_closed(self) -> _Array1D[np.bool_]: ...
+    @property
+    def is_ccw(self) -> _Array1D[np.bool_]: ...
+    @property
+    def has_z(self) -> _Array1D[np.bool_]: ...
+    @property
+    def geom_type(self) -> _Array1D[np.str_]: ...
+    @property
+    def area(self) -> _Array1D[np.float64]: ...
+    @property
+    def length(self) -> _Array1D[np.float64]: ...
+    def count_coordinates(self) -> _Array1D[np.int32]: ...
+    def count_geometries(self) -> _Array1D[np.int32]: ...
+    def count_interior_rings(self) -> _Array1D[np.int32]: ...
+    def get_precision(self) -> _Array1D[np.float64]: ...
+    def get_geometry(self, index: SupportsIndex | ArrayLike) -> _Array1D[np.object_]: ...
+    @property
+    def boundary(self) -> GeometryArray: ...
+    @property
+    def centroid(self) -> GeometryArray: ...
+    def concave_hull(self, ratio: float, allow_holes: bool) -> _Array1D[np.object_]: ...
+    @property
+    def convex_hull(self) -> GeometryArray: ...
+    @property
+    def envelope(self) -> GeometryArray: ...
+    def minimum_rotated_rectangle(self) -> GeometryArray: ...
+    @property
+    def exterior(self) -> GeometryArray: ...
+    def extract_unique_points(self) -> GeometryArray: ...
+    def offset_curve(
+        self,
+        distance: float | ArrayLike,
+        quad_segs: int = 8,
+        join_style: Literal["round", "bevel", "mitre"] = "round",
+        mitre_limit: float = 5.0,
+    ) -> GeometryArray: ...
+    @property
+    def interiors(self) -> _Array1D[np.object_]: ...
+    def remove_repeated_points(self, tolerance: float | ArrayLike = 0.0) -> GeometryArray: ...
+    def representative_point(self) -> GeometryArray: ...
+    def minimum_bounding_circle(self) -> GeometryArray: ...
+    def minimum_bounding_radius(self) -> _Array1D[np.float64]: ...
+    def minimum_clearance(self) -> _Array1D[np.float64]: ...
+    def normalize(self) -> GeometryArray: ...
+    def make_valid(self) -> GeometryArray: ...
+    def reverse(self) -> GeometryArray: ...
+    def segmentize(self, max_segment_length: float | ArrayLike) -> GeometryArray: ...
+    def force_2d(self) -> GeometryArray: ...
+    def force_3d(self, z: float | ArrayLike = 0) -> GeometryArray: ...
+    def transform(
+        self, transformation: Callable[[NDArray[np.float64]], NDArray[np.float64]], include_z: bool = False
+    ) -> GeometryArray: ...
+    def line_merge(self, directed: bool = False) -> GeometryArray: ...
+    def set_precision(
+        self, grid_size: float, mode: Literal["valid_output", "pointwise", "keep_collapsed", 0, 1, 2] = "valid_output"
+    ) -> GeometryArray: ...
+    def covers(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def covered_by(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def contains(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def contains_properly(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def crosses(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def disjoint(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def geom_equals(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def intersects(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def overlaps(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def touches(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def within(self, other: _ArrayOrGeom) -> _Array1D[np.bool_]: ...
+    def dwithin(self, other: _ArrayOrGeom, distance: float) -> _Array1D[np.bool_]: ...
+    def geom_equals_exact(self, other: _ArrayOrGeom, tolerance: float | ArrayLike) -> _Array1D[np.bool_]: ...
+    @deprecated("Use method `geom_equals_exact` instead.")
+    def geom_almost_equals(self, other: _ArrayOrGeom, decimal: float) -> _Array1D[np.bool_]: ...
+    def clip_by_rect(self, xmin: float, ymin: float, xmax: float, ymax: float) -> GeometryArray: ...
+    def difference(self, other: _ArrayOrGeom) -> GeometryArray: ...
+    def intersection(self, other: _ArrayOrGeom) -> GeometryArray: ...
+    def symmetric_difference(self, other: _ArrayOrGeom) -> GeometryArray: ...
+    def union(self, other: _ArrayOrGeom) -> GeometryArray: ...
+    def shortest_line(self, other: _ArrayOrGeom) -> GeometryArray: ...
+    def snap(self, other: _ArrayOrGeom, tolerance: float | ArrayLike) -> GeometryArray: ...
+    def shared_paths(self, other: _ArrayOrGeom) -> GeometryArray: ...
+    def distance(self, other: _ArrayOrGeom) -> _Array1D[np.float64]: ...
+    def hausdorff_distance(self, other: _ArrayOrGeom, **kwargs) -> _Array1D[np.float64]: ...
+    def frechet_distance(self, other: _ArrayOrGeom, **kwargs) -> _Array1D[np.float64]: ...
+    def buffer(self, distance: float | ArrayLike, resolution: int = 16, **kwargs) -> GeometryArray: ...
+    def interpolate(self, distance: float | ArrayLike, normalized: bool = False) -> GeometryArray: ...
+    def simplify(self, tolerance: float | ArrayLike, preserve_topology: bool = True) -> GeometryArray: ...
+    def project(self, other: _ArrayOrGeom, normalized: bool = False) -> _Array1D[np.float64]: ...
+    def relate(self, other: _ArrayOrGeom) -> _Array1D[np.str_]: ...
+    def relate_pattern(self, other: _ArrayOrGeom, pattern: str) -> _Array1D[np.bool_]: ...
+    @deprecated("Use method `union_all` instead.")
+    def unary_union(self) -> BaseGeometry: ...
+    def union_all(self, method: Literal["coverage", "unary"] = "unary") -> BaseGeometry: ...
+    def intersection_all(self) -> BaseGeometry: ...
+    def affine_transform(self, matrix) -> GeometryArray: ...
+    def translate(self, xoff: float = 0.0, yoff: float = 0.0, zoff: float = 0.0) -> GeometryArray: ...
+    def rotate(self, angle: float, origin: _AffinityOrigin = "center", use_radians: bool = False) -> GeometryArray: ...
+    def scale(
+        self, xfact: float = 1.0, yfact: float = 1.0, zfact: float = 1.0, origin: _AffinityOrigin = "center"
+    ) -> GeometryArray: ...
+    def skew(
+        self, xs: float = 0.0, ys: float = 0.0, origin: _AffinityOrigin = "center", use_radians: bool = False
+    ) -> GeometryArray: ...
+    def to_crs(self, crs: _ConvertibleToCRS | None = None, epsg: int | None = None) -> GeometryArray: ...
+    def estimate_utm_crs(self, datum_name: str = "WGS 84") -> CRS: ...
+    @property
+    def x(self) -> _Array1D[np.float64]: ...
+    @property
+    def y(self) -> _Array1D[np.float64]: ...
+    @property
+    def z(self) -> _Array1D[np.float64]: ...
+    @property
+    def bounds(self) -> _Array2D[np.float64]: ...
+    @property
+    def total_bounds(self) -> _Array1D[np.float64]: ...
+    @property
+    def size(self) -> int: ...
+    @property
+    def shape(self) -> tuple[int]: ...  # Always 1-D, this is not a mistake
+    @property
+    def ndim(self) -> Literal[1]: ...
+    def copy(self, *args: Unused, **kwargs: Unused) -> GeometryArray: ...
+    def take(
+        self,
+        indices: Sequence[SupportsIndex] | NDArray[np.integer[Any]],  # np.integer[Any] because precision is not important
+        allow_fill: bool = False,
+        fill_value: Geometry | None = None,
+    ) -> GeometryArray: ...
+    def fillna(
+        self,
+        value: Geometry | GeometryArray | None = None,
+        method: Literal["backfill", "bfill", "pad", "ffill"] | None = None,
+        limit: int | None = None,
+        copy: bool = True,
+    ) -> GeometryArray: ...
+    @overload
+    def astype(self, dtype: GeometryDtype, copy: bool = True) -> GeometryArray: ...
+    @overload
+    def astype(self, dtype: ExtensionDtype | Literal["string"], copy: bool = True) -> ExtensionArray: ...  # type: ignore[overload-overlap]
+    @overload
+    def astype(self, dtype: DTypeLike, copy: bool = True) -> _Array1D[Incomplete]: ...
+    def isna(self) -> _Array1D[np.bool_]: ...
+    def value_counts(self, dropna: bool = True) -> pd.Series[int]: ...
+    def unique(self) -> GeometryArray: ...
+    @property
+    def nbytes(self) -> int: ...
+    def shift(self, periods: int = 1, fill_value: Geometry | None = None) -> GeometryArray: ...  # type: ignore[override]
+    def argmin(self, skipna: bool = True) -> NoReturn: ...
+    def argmax(self, skipna: bool = True) -> NoReturn: ...
+    def __array__(self, dtype: DTypeLike | None = None, copy: bool | None = None) -> _Array1D[np.object_]: ...
+    def __eq__(self, other: object) -> _Array1D[np.bool_]: ...  # type: ignore[override]
+    def __ne__(self, other: object) -> _Array1D[np.bool_]: ...  # type: ignore[override]
+    def __contains__(self, item: object) -> np.bool_: ...
+
+def transform(
+    data: NDArray[np.object_], func: Callable[[NDArray[np.float64], NDArray[np.float64]], NDArray[np.float64]]
+) -> NDArray[np.object_]: ...