Skip to content

Commit 122a100

Browse files
committed
Improve _geoarrow module
1 parent 5f20ec1 commit 122a100

File tree

1 file changed

+30
-11
lines changed

1 file changed

+30
-11
lines changed

stubs/geopandas/geopandas/io/_geoarrow.pyi

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,46 @@
11
from _typeshed import Incomplete
2-
from typing import Literal
3-
from typing_extensions import TypeAlias
2+
from collections.abc import Mapping
3+
from typing import (
4+
# pyarrow types returned as Any to avoid depending on pyarrow (40 MB) in stubs
5+
Any as _PAArray,
6+
Any as _PAField,
7+
Any as _PATable,
8+
Literal,
9+
Protocol,
10+
type_check_only,
11+
)
12+
from typing_extensions import CapsuleType, TypeAlias
413

514
import numpy as np
615
from numpy.typing import NDArray
716

817
from ..array import GeometryArray
918
from ..geodataframe import GeoDataFrame
1019

11-
_PATable: TypeAlias = Incomplete
12-
_PAField: TypeAlias = Incomplete
13-
_PAArray: TypeAlias = Incomplete
14-
1520
# Literal for language server completions and str because runtime normalizes to lowercase
1621
_GeomEncoding: TypeAlias = Literal["WKB", "geoarrow"] | str # noqa: Y051
1722

23+
@type_check_only
24+
class _PyarrowTableLike(Protocol):
25+
def __arrow_c_stream__(self, requested_schema=None) -> CapsuleType: ...
26+
27+
@type_check_only
28+
class _PyarrowFieldLike(Protocol):
29+
def __arrow_c_schema__(self) -> CapsuleType: ...
30+
31+
@type_check_only
32+
class _PyarrowArrayLike(Protocol):
33+
def __arrow_c_array__(self) -> tuple[CapsuleType, CapsuleType]: ...
34+
1835
GEOARROW_ENCODINGS: list[str]
1936

2037
class ArrowTable:
21-
def __init__(self, pa_table: _PATable) -> None: ...
22-
def __arrow_c_stream__(self, requested_schema=None): ...
38+
def __init__(self, pa_table: _PyarrowTableLike) -> None: ...
39+
def __arrow_c_stream__(self, requested_schema=None) -> CapsuleType: ...
2340

2441
class GeoArrowArray:
25-
def __init__(self, pa_field: _PAField, pa_array: _PAArray) -> None: ...
26-
def __arrow_c_array__(self, requested_schema=None) -> tuple[Incomplete, Incomplete]: ...
42+
def __init__(self, pa_field: _PyarrowFieldLike, pa_array: _PyarrowArrayLike) -> None: ...
43+
def __arrow_c_array__(self, requested_schema=None) -> tuple[CapsuleType, CapsuleType]: ...
2744

2845
def geopandas_to_arrow(
2946
df: GeoDataFrame,
@@ -43,6 +60,8 @@ def construct_geometry_array(
4360
crs: str | None = None,
4461
interleaved: bool = True,
4562
) -> tuple[_PAField, _PAArray]: ...
46-
def arrow_to_geopandas(table, geometry: str | None = None) -> GeoDataFrame: ...
63+
def arrow_to_geopandas(
64+
table, geometry: str | None = None, to_pandas_kwargs: Mapping[str, Incomplete] | None = None
65+
) -> GeoDataFrame: ...
4766
def arrow_to_geometry_array(arr) -> GeometryArray: ...
4867
def construct_shapely_array(arr: _PAArray, extension_name: str) -> NDArray[np.object_]: ...

0 commit comments

Comments
 (0)