rok commented on code in PR #55264: URL: https://github.com/apache/arrow/pull/55264#discussion_r4515157324
########## python/pyarrow-stubs/pyarrow/lib.pyi: ########## @@ -0,0 +1,36 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Stub for the pyarrow.lib C extension module. Re-exports the modules that +# have been annotated so far; everything else falls back to Any via +# __getattr__ until it gets annotated in subsequent PRs. + +from typing import Any + +from pyarrow.__lib.types import * # noqa: F401,F403 +from pyarrow.__lib.array import * # noqa: F401,F403 +from pyarrow.__lib.builder import * # noqa: F401,F403 +from pyarrow.__lib.config import * # noqa: F401,F403 +from pyarrow.__lib.device import * # noqa: F401,F403 +from pyarrow.__lib.error import * # noqa: F401,F403 +from pyarrow.__lib.io import * # noqa: F401,F403 +from pyarrow.__lib.memory import * # noqa: F401,F403 +from pyarrow.__lib.scalar import * # noqa: F401,F403 +from pyarrow.__lib.table import * # noqa: F401,F403 +from pyarrow.__lib.tensor import * # noqa: F401,F403 + Review Comment: ```suggestion class MonthDayNano(tuple): months: int days: int nanoseconds: int def __new__( cls, sequence: tuple[int, int, int] | list[int] = ..., ) -> MonthDayNano: ... ``` ########## python/pyarrow-stubs/pyarrow/__lib/table.pyi: ########## @@ -0,0 +1,613 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import Self, TypeAlias + +from collections.abc import Iterable, Iterator, Sequence, Mapping +from typing import Any, ClassVar, Literal +from typing_extensions import CapsuleType, deprecated +import builtins + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] + +from numpy.typing import NDArray # type: ignore[import-not-found] +from pyarrow._compute import ( + CastOptions, + CountOptions, + FunctionOptions, + ScalarAggregateOptions, + TDigestOptions, + VarianceOptions, +) +from pyarrow._stubs_typing import ( + Indices, + Mask, + MapsAsPydicts, + NullEncoding, + NullPlacement, + NullSelectionBehavior, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsArrowStream, + SupportsPyArrowArray, +) +from pyarrow.compute import Expression +from pyarrow.interchange.dataframe import _PyArrowDataFrame +from pyarrow.lib import Device, Field, MemoryManager, MemoryPool, Schema + +from .array import Array, StructArray, _PandasConvertible +from .device import DeviceAllocationType +from .io import Buffer +from .ipc import RecordBatchReader +from .scalar import Int64Scalar, Scalar +from .tensor import Tensor +from .types import DataType + +ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray + +_Aggregation: TypeAlias = Literal[ + "all", + "any", + "approximate_median", + "count", + "count_all", + "count_distinct", + "distinct", + "first", + "first_last", + "kurtosis", + "last", + "list", + "max", + "mean", + "min", + "min_max", + "one", + "pivot_wider", + "product", + "skew", + "stddev", + "sum", + "tdigest", + "variance", +] +_AggregationPrefixed: TypeAlias = Literal[ + "hash_all", + "hash_any", + "hash_approximate_median", + "hash_count", + "hash_count_all", + "hash_count_distinct", + "hash_distinct", + "hash_first", + "hash_first_last", + "hash_kurtosis", + "hash_last", + "hash_list", + "hash_max", + "hash_mean", + "hash_min", + "hash_min_max", + "hash_one", + "hash_pivot_wider", + "hash_product", + "hash_skew", + "hash_stddev", + "hash_sum", + "hash_tdigest", + "hash_variance", +] +Aggregation: TypeAlias = _Aggregation | _AggregationPrefixed | str +AggregateOptions: TypeAlias = ( + ScalarAggregateOptions + | CountOptions + | TDigestOptions + | VarianceOptions + | FunctionOptions +) + +ChunkedArrayConstructorInput: TypeAlias = ( Review Comment: Will this be useful outside of `table.pyi` and should we move it to `_stubs_typing.pyi`? ########## docs/source/developers/python/development.rst: ########## @@ -169,6 +181,48 @@ in downstream projects and for users' IDEs. Note: ``py.typed`` marker file in the ``pyarrow/`` directory indicates to type checkers that PyArrow supports type checking according to :pep:`561`. +Testing Type Stubs +------------------ + +The tests in ``pyarrow-stubs/tests/`` pin down the annotated types of public +APIs. Each check pairs :func:`typing.assert_type`, which every type checker +verifies statically, with a runtime ``isinstance`` check: + +.. code-block:: python + + import pyarrow.lib as lib + + from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + + def test_list_type() -> None: + lst = lib.list_(lib.int32()) + check(assert_type(lst, lib.ListType), lib.ListType) + check(assert_type(lst.value_type, lib.DataType), lib.DataType) + + if TYPE_CHECKING_INVALID_USAGE: + lib.list_(1) # type: ignore[call-overload] + +The type checkers check these files as part of the regular runs above, and +pytest runs them against the installed PyArrow to make sure the stubs agree +with what PyArrow actually returns: + +.. code-block:: + + $ cd arrow/python + $ pytest pyarrow-stubs/tests + +Run ``pytest`` directly and pass it the directory: ``python -m pytest`` puts +``arrow/python`` on ``sys.path``, so the source tree would be imported instead +of the installed PyArrow. Review Comment: I think this is Claude-ish, sorry, let's do regular `pytest`. ########## python/pyarrow-stubs/pyarrow/__lib/types.pyi: ########## @@ -0,0 +1,700 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Any, Literal, TypeAlias, overload + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] +from typing_extensions import CapsuleType, Self, deprecated + +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow.lib import ( + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + Table, +) + +class _Weakrefable: ... +class _Metadata(_Weakrefable): ... + +class DataType(_Weakrefable): + def field(self, i: int) -> Field: ... + @property + def id(self) -> int: ... + @property + def bit_width(self) -> int: ... + @property + def byte_width(self) -> int: ... + @property + def num_fields(self) -> int: ... + @property + def num_buffers(self) -> int: ... + @property + def has_variadic_buffers(self) -> bool: ... + def __hash__(self) -> int: ... + def equals( + self, other: DataType | str, *, check_metadata: bool = False + ) -> bool: ... + def to_pandas_dtype(self) -> _PandasDtype: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> DataType: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> DataType: ... + +_PandasDtype: TypeAlias = ( + type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype +) + +class NullType(DataType): ... +class BoolType(DataType): ... +class UInt8Type(DataType): ... +class Int8Type(DataType): ... +class UInt16Type(DataType): ... +class Int16Type(DataType): ... +class UInt32Type(DataType): ... +class Int32Type(DataType): ... +class UInt64Type(DataType): ... +class Int64Type(DataType): ... +class Float16Type(DataType): ... +class Float32Type(DataType): ... +class Float64Type(DataType): ... +class Date32Type(DataType): ... +class Date64Type(DataType): ... +class MonthDayNanoIntervalType(DataType): ... +class StringType(DataType): ... +class LargeStringType(DataType): ... +class StringViewType(DataType): ... +class BinaryType(DataType): ... +class LargeBinaryType(DataType): ... +class BinaryViewType(DataType): ... + +_Time32Unit: TypeAlias = Literal["s", "ms"] +_Time64Unit: TypeAlias = Literal["us", "ns"] + +class TimestampType(DataType): + @property + def unit(self) -> TimeUnit: ... + @property + def tz(self) -> str | None: ... + +class Time32Type(DataType): + @property + def unit(self) -> _Time32Unit: ... + +class Time64Type(DataType): + @property + def unit(self) -> _Time64Unit: ... + +class DurationType(DataType): + @property + def unit(self) -> TimeUnit: ... + +class FixedSizeBinaryType(DataType): ... + +class Decimal32Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal64Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal128Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal256Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class ListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class ListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class FixedSizeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + @property + def list_size(self) -> int: ... + +class DictionaryMemo(_Weakrefable): ... + +_IndexType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type +) + +class DictionaryType(DataType): + @property + def ordered(self) -> bool: ... + @property + def index_type(self) -> _IndexType: ... + @property + def value_type(self) -> DataType: ... + +class MapType(DataType): + @property + def key_field(self) -> Field: ... + @property + def key_type(self) -> DataType: ... + @property + def item_field(self) -> Field: ... + @property + def item_type(self) -> DataType: ... + @property + def keys_sorted(self) -> bool: ... + +class StructType(DataType): + def get_field_index(self, name: str | bytes) -> int: ... + def field(self, i: int | str | bytes) -> Field: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + + __getitem__ = field + @property + def names(self) -> list[str]: ... + @property + def fields(self) -> list[Field]: ... + +class UnionType(DataType): + @property + def mode(self) -> Literal["sparse", "dense"]: ... + @property + def type_codes(self) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + def field(self, i: int) -> Field: ... + + __getitem__ = field + +class SparseUnionType(UnionType): + @property + def mode(self) -> Literal["sparse"]: ... + +class DenseUnionType(UnionType): + @property + def mode(self) -> Literal["dense"]: ... + +_RunEndType: TypeAlias = Int16Type | Int32Type | Int64Type + +class RunEndEncodedType(DataType): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> DataType: ... + +class BaseExtensionType(DataType): + def __arrow_ext_class__(self) -> type[ExtensionArray]: ... + def __arrow_ext_scalar_class__(self) -> type[ExtensionScalar]: ... + @property + def extension_name(self) -> str: ... + @property + def storage_type(self) -> DataType: ... + def wrap_array( + self, storage: Array | ChunkedArray + ) -> ExtensionArray | ChunkedArray: ... + +class ExtensionType(BaseExtensionType): + def __init__(self, storage_type: DataType, extension_name: str) -> None: ... + def __arrow_ext_serialize__(self) -> bytes: ... + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: DataType, serialized: bytes + ) -> Self: ... + +class FixedShapeTensorType(BaseExtensionType): + @property + def value_type(self) -> DataType: ... + @property + def shape(self) -> list[int]: ... + @property + def dim_names(self) -> list[str] | None: ... + @property + def permutation(self) -> list[int] | None: ... + +class Bool8Type(BaseExtensionType): ... +class UuidType(BaseExtensionType): ... +class JsonType(BaseExtensionType): ... + +class OpaqueType(BaseExtensionType): + @property + def type_name(self) -> str: ... + @property + def vendor_name(self) -> str: ... + +class UnknownExtensionType(ExtensionType): + def __init__(self, storage_type: DataType, serialized: bytes) -> None: ... + +def register_extension_type(ext_type: BaseExtensionType) -> None: ... +def unregister_extension_type(type_name: str) -> None: ... + +_FieldName: TypeAlias = str | bytes +_StrOrBytes: TypeAlias = str | bytes +_MetadataMapping: TypeAlias = ( + Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] +) +_MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] +_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable | None +_DataTypeAlias: TypeAlias = Literal[ + "null", + "bool", + "boolean", + "i1", + "int8", + "i2", + "int16", + "i4", + "int32", + "i8", + "int64", + "u1", + "uint8", + "u2", + "uint16", + "u4", + "uint32", + "u8", + "uint64", + "f2", + "halffloat", + "float16", + "f4", + "float", + "float32", + "f8", + "double", + "float64", + "string", + "str", + "utf8", + "binary", + "large_string", + "large_str", + "large_utf8", + "large_binary", + "binary_view", + "string_view", + "date32", + "date64", + "date32[day]", + "date64[ms]", + "time32[s]", + "time32[ms]", + "time64[us]", + "time64[ns]", + "timestamp[s]", + "timestamp[ms]", + "timestamp[us]", + "timestamp[ns]", + "duration[s]", + "duration[ms]", + "duration[us]", + "duration[ns]", + "month_day_nano_interval", +] +_DataTypeAliasInput: TypeAlias = _DataTypeAlias | str +_DataTypeLike: TypeAlias = DataType | _DataTypeAliasInput + +class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): + def __init__( + self, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, Review Comment: I'd say yes. ########## python/pyarrow-stubs/pyarrow/__lib/types.pyi: ########## @@ -0,0 +1,700 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Any, Literal, TypeAlias, overload + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] +from typing_extensions import CapsuleType, Self, deprecated + +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow.lib import ( + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + Table, +) + +class _Weakrefable: ... +class _Metadata(_Weakrefable): ... + +class DataType(_Weakrefable): + def field(self, i: int) -> Field: ... + @property + def id(self) -> int: ... + @property + def bit_width(self) -> int: ... + @property + def byte_width(self) -> int: ... + @property + def num_fields(self) -> int: ... + @property + def num_buffers(self) -> int: ... + @property + def has_variadic_buffers(self) -> bool: ... + def __hash__(self) -> int: ... + def equals( + self, other: DataType | str, *, check_metadata: bool = False + ) -> bool: ... + def to_pandas_dtype(self) -> _PandasDtype: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> DataType: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> DataType: ... + +_PandasDtype: TypeAlias = ( + type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype +) + +class NullType(DataType): ... +class BoolType(DataType): ... +class UInt8Type(DataType): ... +class Int8Type(DataType): ... +class UInt16Type(DataType): ... +class Int16Type(DataType): ... +class UInt32Type(DataType): ... +class Int32Type(DataType): ... +class UInt64Type(DataType): ... +class Int64Type(DataType): ... +class Float16Type(DataType): ... +class Float32Type(DataType): ... +class Float64Type(DataType): ... +class Date32Type(DataType): ... +class Date64Type(DataType): ... +class MonthDayNanoIntervalType(DataType): ... +class StringType(DataType): ... +class LargeStringType(DataType): ... +class StringViewType(DataType): ... +class BinaryType(DataType): ... +class LargeBinaryType(DataType): ... +class BinaryViewType(DataType): ... + +_Time32Unit: TypeAlias = Literal["s", "ms"] +_Time64Unit: TypeAlias = Literal["us", "ns"] + +class TimestampType(DataType): + @property + def unit(self) -> TimeUnit: ... + @property + def tz(self) -> str | None: ... + +class Time32Type(DataType): + @property + def unit(self) -> _Time32Unit: ... + +class Time64Type(DataType): + @property + def unit(self) -> _Time64Unit: ... + +class DurationType(DataType): + @property + def unit(self) -> TimeUnit: ... + +class FixedSizeBinaryType(DataType): ... + +class Decimal32Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal64Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal128Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal256Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class ListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class ListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class FixedSizeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + @property + def list_size(self) -> int: ... + +class DictionaryMemo(_Weakrefable): ... + +_IndexType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type +) + +class DictionaryType(DataType): + @property + def ordered(self) -> bool: ... + @property + def index_type(self) -> _IndexType: ... + @property + def value_type(self) -> DataType: ... + +class MapType(DataType): + @property + def key_field(self) -> Field: ... + @property + def key_type(self) -> DataType: ... + @property + def item_field(self) -> Field: ... + @property + def item_type(self) -> DataType: ... + @property + def keys_sorted(self) -> bool: ... + +class StructType(DataType): + def get_field_index(self, name: str | bytes) -> int: ... + def field(self, i: int | str | bytes) -> Field: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + + __getitem__ = field + @property + def names(self) -> list[str]: ... + @property + def fields(self) -> list[Field]: ... + +class UnionType(DataType): + @property + def mode(self) -> Literal["sparse", "dense"]: ... + @property + def type_codes(self) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + def field(self, i: int) -> Field: ... + + __getitem__ = field + +class SparseUnionType(UnionType): + @property + def mode(self) -> Literal["sparse"]: ... + +class DenseUnionType(UnionType): + @property + def mode(self) -> Literal["dense"]: ... + +_RunEndType: TypeAlias = Int16Type | Int32Type | Int64Type + +class RunEndEncodedType(DataType): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> DataType: ... + +class BaseExtensionType(DataType): + def __arrow_ext_class__(self) -> type[ExtensionArray]: ... + def __arrow_ext_scalar_class__(self) -> type[ExtensionScalar]: ... + @property + def extension_name(self) -> str: ... + @property + def storage_type(self) -> DataType: ... + def wrap_array( + self, storage: Array | ChunkedArray + ) -> ExtensionArray | ChunkedArray: ... + +class ExtensionType(BaseExtensionType): + def __init__(self, storage_type: DataType, extension_name: str) -> None: ... + def __arrow_ext_serialize__(self) -> bytes: ... + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: DataType, serialized: bytes + ) -> Self: ... + +class FixedShapeTensorType(BaseExtensionType): + @property + def value_type(self) -> DataType: ... + @property + def shape(self) -> list[int]: ... + @property + def dim_names(self) -> list[str] | None: ... + @property + def permutation(self) -> list[int] | None: ... + +class Bool8Type(BaseExtensionType): ... +class UuidType(BaseExtensionType): ... +class JsonType(BaseExtensionType): ... + +class OpaqueType(BaseExtensionType): + @property + def type_name(self) -> str: ... + @property + def vendor_name(self) -> str: ... + +class UnknownExtensionType(ExtensionType): + def __init__(self, storage_type: DataType, serialized: bytes) -> None: ... + +def register_extension_type(ext_type: BaseExtensionType) -> None: ... +def unregister_extension_type(type_name: str) -> None: ... + +_FieldName: TypeAlias = str | bytes +_StrOrBytes: TypeAlias = str | bytes +_MetadataMapping: TypeAlias = ( + Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] +) +_MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] +_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable | None +_DataTypeAlias: TypeAlias = Literal[ + "null", + "bool", + "boolean", + "i1", + "int8", + "i2", + "int16", + "i4", + "int32", + "i8", + "int64", + "u1", + "uint8", + "u2", + "uint16", + "u4", + "uint32", + "u8", + "uint64", + "f2", + "halffloat", + "float16", + "f4", + "float", + "float32", + "f8", + "double", + "float64", + "string", + "str", + "utf8", + "binary", + "large_string", + "large_str", + "large_utf8", + "large_binary", + "binary_view", + "string_view", + "date32", + "date64", + "date32[day]", + "date64[ms]", + "time32[s]", + "time32[ms]", + "time64[us]", + "time64[ns]", + "timestamp[s]", + "timestamp[ms]", + "timestamp[us]", + "timestamp[ns]", + "duration[s]", + "duration[ms]", + "duration[us]", + "duration[ns]", + "month_day_nano_interval", +] +_DataTypeAliasInput: TypeAlias = _DataTypeAlias | str +_DataTypeLike: TypeAlias = DataType | _DataTypeAliasInput + +class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): + def __init__( + self, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, + **kwargs: str | bytes, + ) -> None: ... + def equals(self, other: KeyValueMetadata) -> bool: ... + def __len__(self) -> int: ... + def __contains__(self, /, __key: object) -> bool: ... + def __getitem__(self, /, __key: str | bytes) -> bytes: ... + def __iter__(self) -> Iterator[bytes]: ... + def key(self, i: int) -> bytes: ... + def value(self, i: int) -> bytes: ... + # The runtime yields generators instead of Mapping views + def keys(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def values(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def items(self) -> Iterator[tuple[bytes, bytes]]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def get_all(self, key: str | bytes) -> list[bytes]: ... + def to_dict(self) -> dict[bytes, bytes]: ... + +class Field(_Weakrefable): + def equals(self, other: Field, check_metadata: bool = False) -> bool: ... + def __hash__(self) -> int: ... + @property + def nullable(self) -> bool: ... + @property + def name(self) -> str: ... + @property + def metadata(self) -> dict[bytes, bytes] | None: ... + @property + def type(self) -> DataType: ... + def with_metadata( + self, + metadata: _KeyValueMetadataInput | KeyValueMetadata, + ) -> Field: ... + def remove_metadata(self) -> Field: ... + def with_type(self, new_type: DataType) -> Field: ... + def with_name(self, name: str | bytes) -> Field: ... + def with_nullable(self, nullable: bool) -> Field: ... + def flatten(self) -> list[Field]: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> Field: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Field: ... + +_FieldTuple: TypeAlias = ( + tuple[_FieldName, _DataTypeLike] + | tuple[_FieldName, _DataTypeLike, bool | None] + | tuple[ + _FieldName, + _DataTypeLike, + bool | None, + _KeyValueMetadataInput | KeyValueMetadata, + ] +) +_StructFieldsInput: TypeAlias = ( + Iterable[Field | _FieldTuple] | Mapping[_FieldName, _DataTypeLike] +) +_SchemaFieldsInput: TypeAlias = _StructFieldsInput + +class Schema(_Weakrefable): + def __len__(self) -> int: ... + def __getitem__(self, key: int) -> Field: ... + + _field = __getitem__ + def __iter__(self) -> Iterator[Field]: ... + def __hash__(self) -> int: ... + def __sizeof__(self) -> int: ... + @property + def pandas_metadata(self) -> dict[str, Any] | None: ... + @property + def names(self) -> list[str]: ... + @property + def types(self) -> list[DataType]: ... + @property + def metadata(self) -> dict[bytes, bytes] | None: ... + def empty_table(self) -> Table: ... + def equals(self, other: Schema, check_metadata: bool = False) -> bool: ... + @classmethod + def from_pandas( + cls, df: pd.DataFrame, preserve_index: bool | None = None + ) -> Schema: ... + def field(self, i: int | str | bytes) -> Field: ... + @deprecated("Use 'field' instead") + def field_by_name(self, name: str | bytes) -> Field | None: ... + def get_field_index(self, name: str | bytes) -> int: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def append(self, field: Field) -> Schema: ... + def insert(self, i: int, field: Field) -> Schema: ... + def remove(self, i: int) -> Schema: ... + def set(self, i: int, field: Field) -> Schema: ... + @deprecated("Use 'with_metadata' instead") + def add_metadata( + self, metadata: _KeyValueMetadataInput | KeyValueMetadata + ) -> Schema: ... + def with_metadata( + self, metadata: _KeyValueMetadataInput | KeyValueMetadata + ) -> Schema: ... + def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... + def remove_metadata(self) -> Schema: ... + def to_string( + self, + truncate_metadata: bool = True, + show_field_metadata: bool = True, + show_schema_metadata: bool = True, + element_size_limit: int = 100, + ) -> str: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> Schema: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Schema: ... + +def unify_schemas( + schemas: Iterable[Schema], + *, + promote_options: Literal["default", "permissive"] = "default", +) -> Schema: ... +@overload +def field( + name: SupportsArrowSchema, + type: None = None, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata = None, +) -> Field: ... +@overload +def field( + name: _FieldName, + type: _DataTypeLike, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata = None, +) -> Field: ... +def null() -> NullType: ... +def bool_() -> BoolType: ... +def uint8() -> UInt8Type: ... +def int8() -> Int8Type: ... +def uint16() -> UInt16Type: ... +def int16() -> Int16Type: ... +def uint32() -> UInt32Type: ... +def int32() -> Int32Type: ... +def int64() -> Int64Type: ... +def uint64() -> UInt64Type: ... +def timestamp(unit: TimeUnit | str, tz: str | None = None) -> TimestampType: ... +def time32(unit: _Time32Unit | str) -> Time32Type: ... +def time64(unit: _Time64Unit | str) -> Time64Type: ... +def duration(unit: TimeUnit | str) -> DurationType: ... +def month_day_nano_interval() -> MonthDayNanoIntervalType: ... +def date32() -> Date32Type: ... +def date64() -> Date64Type: ... +def float16() -> Float16Type: ... +def float32() -> Float32Type: ... +def float64() -> Float64Type: ... +def decimal32(precision: int, scale: int = 0) -> Decimal32Type: ... +def decimal64(precision: int, scale: int = 0) -> Decimal64Type: ... +def decimal128(precision: int, scale: int = 0) -> Decimal128Type: ... +def decimal256(precision: int, scale: int = 0) -> Decimal256Type: ... +def string() -> StringType: ... + +utf8 = string + +def binary( + length: Literal[-1] | int = ..., +) -> BinaryType | FixedSizeBinaryType: ... Review Comment: I think adding overloads here are ok, but we should be careful they don't explode on us. Maybe we should have a heuristic - e.g. if we need more than 3 overloads we should discuss before adding? I've looked at old discussions on other threads and sentiment is generally positive on overloads, concern is about them growing uncontrollably. ########## python/pyarrow-stubs/tests/test_types.py: ########## @@ -0,0 +1,229 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# mypy: warn-unused-ignores + +from collections.abc import Hashable, Iterator +from typing import TYPE_CHECKING, Literal + +import pyarrow.lib as lib +import pyarrow.types as types +from typing_extensions import CapsuleType + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + +if TYPE_CHECKING: + # The stubs have a class per primitive type, but at runtime the + # factories return plain DataType instances, so these can only be + # referred to from type checking code. + from pyarrow.lib import ( # noqa: F401 + BoolType, Review Comment: Agreed, this should be changed. TODO(rok). ########## python/pyarrow-stubs/tests/test_array.py: ########## @@ -0,0 +1,136 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# mypy: warn-unused-ignores + +from typing import Any, cast + +import pyarrow.lib as lib +from typing_extensions import CapsuleType + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + + +def test_array() -> None: + # pa.array() isn't annotated with a specific return type yet, so tell + # the type checker which class to expect. + arr = check(cast(lib.Int64Array, lib.array([1, 2, None])), lib.Int64Array) Review Comment: I think we should switch to `pa.*`, buut we have `def __getattr__(name: str) -> Any: ...` in `__init__.pyi` at the moment. How about we flip this as we merge the last PR in the series? ########## python/pyarrow-stubs/tests/test_types.py: ########## @@ -0,0 +1,229 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# mypy: warn-unused-ignores + +from collections.abc import Hashable, Iterator +from typing import TYPE_CHECKING, Literal + +import pyarrow.lib as lib +import pyarrow.types as types +from typing_extensions import CapsuleType + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + +if TYPE_CHECKING: + # The stubs have a class per primitive type, but at runtime the + # factories return plain DataType instances, so these can only be + # referred to from type checking code. + from pyarrow.lib import ( # noqa: F401 + BoolType, + Int8Type, + Int16Type, + Int32Type, + Int64Type, + StringType, + UInt8Type, + UInt16Type, + UInt32Type, + UInt64Type, + ) + + IntegerType = ( + UInt8Type | Int8Type | UInt16Type | Int16Type + | UInt32Type | Int32Type | UInt64Type | Int64Type + ) + + +def _erase(t: lib.DataType) -> lib.DataType: + # Hide the precise type, so we can test narrowing by the is_* functions. + return t + + +def test_primitive_factories() -> None: + check(assert_type(lib.bool_(), "BoolType"), lib.DataType) + check(assert_type(lib.int32(), "Int32Type"), lib.DataType) + check(assert_type(lib.string(), "StringType"), lib.DataType) + check(assert_type(lib.utf8(), "StringType"), lib.DataType) + + +def test_temporal_types() -> None: + ts = lib.timestamp("ms", tz="UTC") + check(assert_type(ts, lib.TimestampType), lib.TimestampType) + check(assert_type(ts.unit, Literal["s", "ms", "us", "ns"]), str) + check(assert_type(ts.tz, str | None), str) + + check(assert_type(lib.time32("s").unit, Literal["s", "ms"]), str) + check(assert_type(lib.time64("ns").unit, Literal["us", "ns"]), str) + check(assert_type(lib.duration("us"), lib.DurationType), lib.DurationType) + + +def test_decimal_types() -> None: + dec = lib.decimal128(10, 2) + check(assert_type(dec, lib.Decimal128Type), lib.Decimal128Type) + check(assert_type(dec.precision, int), int) + check(assert_type(dec.scale, int), int) + + +def test_list_types() -> None: + lst = lib.list_(lib.int32()) + check(assert_type(lst, lib.ListType), lib.ListType) + check(assert_type(lst.value_type, lib.DataType), lib.DataType) + check(assert_type(lst.value_field, lib.Field), lib.Field) + + # With an explicit list_size we can't tell statically whether it is -1. + fixed = lib.list_(lib.int32(), 2) + check(assert_type(fixed, lib.ListType | lib.FixedSizeListType), Review Comment: So even with overload we can't narrow this? ########## python/pyarrow-stubs/tests/test_types.py: ########## @@ -0,0 +1,229 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# mypy: warn-unused-ignores + +from collections.abc import Hashable, Iterator +from typing import TYPE_CHECKING, Literal + +import pyarrow.lib as lib +import pyarrow.types as types +from typing_extensions import CapsuleType + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + +if TYPE_CHECKING: + # The stubs have a class per primitive type, but at runtime the + # factories return plain DataType instances, so these can only be + # referred to from type checking code. + from pyarrow.lib import ( # noqa: F401 + BoolType, + Int8Type, + Int16Type, + Int32Type, + Int64Type, + StringType, + UInt8Type, + UInt16Type, + UInt32Type, + UInt64Type, + ) + + IntegerType = ( + UInt8Type | Int8Type | UInt16Type | Int16Type + | UInt32Type | Int32Type | UInt64Type | Int64Type + ) + + +def _erase(t: lib.DataType) -> lib.DataType: + # Hide the precise type, so we can test narrowing by the is_* functions. + return t + + +def test_primitive_factories() -> None: + check(assert_type(lib.bool_(), "BoolType"), lib.DataType) + check(assert_type(lib.int32(), "Int32Type"), lib.DataType) + check(assert_type(lib.string(), "StringType"), lib.DataType) + check(assert_type(lib.utf8(), "StringType"), lib.DataType) + + +def test_temporal_types() -> None: + ts = lib.timestamp("ms", tz="UTC") + check(assert_type(ts, lib.TimestampType), lib.TimestampType) + check(assert_type(ts.unit, Literal["s", "ms", "us", "ns"]), str) + check(assert_type(ts.tz, str | None), str) + + check(assert_type(lib.time32("s").unit, Literal["s", "ms"]), str) + check(assert_type(lib.time64("ns").unit, Literal["us", "ns"]), str) + check(assert_type(lib.duration("us"), lib.DurationType), lib.DurationType) + + +def test_decimal_types() -> None: + dec = lib.decimal128(10, 2) + check(assert_type(dec, lib.Decimal128Type), lib.Decimal128Type) + check(assert_type(dec.precision, int), int) + check(assert_type(dec.scale, int), int) + + +def test_list_types() -> None: + lst = lib.list_(lib.int32()) + check(assert_type(lst, lib.ListType), lib.ListType) + check(assert_type(lst.value_type, lib.DataType), lib.DataType) + check(assert_type(lst.value_field, lib.Field), lib.Field) + + # With an explicit list_size we can't tell statically whether it is -1. + fixed = lib.list_(lib.int32(), 2) + check(assert_type(fixed, lib.ListType | lib.FixedSizeListType), + lib.FixedSizeListType) + + check(assert_type(lib.large_list(lib.int32()), lib.LargeListType), + lib.LargeListType) + check(assert_type(lib.list_view(lib.int32()), lib.ListViewType), + lib.ListViewType) + + +def test_dictionary_type() -> None: + dict_type = lib.dictionary(lib.int8(), lib.string()) + check(assert_type(dict_type, lib.DictionaryType), lib.DictionaryType) + check(assert_type(dict_type.ordered, bool), bool) + check(assert_type(dict_type.value_type, lib.DataType), lib.DataType) + check(assert_type(dict_type.index_type, "IntegerType"), lib.DataType) + + +def test_map_and_run_end_encoded_types() -> None: + map_type = lib.map_(lib.string(), lib.int64()) + check(assert_type(map_type, lib.MapType), lib.MapType) + check(assert_type(map_type.key_type, lib.DataType), lib.DataType) + check(assert_type(map_type.keys_sorted, bool), bool) + + ree = lib.run_end_encoded(lib.int16(), lib.string()) + check(assert_type(ree, lib.RunEndEncodedType), lib.RunEndEncodedType) + check(assert_type(ree.run_end_type, "Int16Type | Int32Type | Int64Type"), + lib.DataType) + + +def test_types_enum() -> None: + assert lib.int64().id == types.TypesEnum.INT64 + assert lib.list_(lib.int8()).id == types.TypesEnum.LIST + assert (lib.month_day_nano_interval().id Review Comment: Let's remove. TODO(rok). ########## python/pyarrow/tests/test_array.py: ########## @@ -323,11 +327,11 @@ def test_asarray(): arr = pa.array(range(4)) - # The iterator interface gives back an array of Int64Value's + # The iterator interface gives back an array of Int64Scalar's np_arr = np.asarray([_ for _ in arr]) assert np_arr.tolist() == [0, 1, 2, 3] assert np_arr.dtype == np.dtype('O') - assert isinstance(np_arr[0], pa.lib.Int64Value) + assert isinstance(np_arr[0], pa.Int64Scalar) Review Comment: Test is decorated ```python @pytest.mark.nopandas @pytest.mark.pandas ``` Does this mean it's always ran or never? :D ########## python/pyarrow/tests/test_array.py: ########## @@ -18,20 +18,24 @@ from collections.abc import Iterable import datetime import decimal -import hypothesis as h -import hypothesis.strategies as st +import hypothesis as h # type: ignore[import-not-found] +import hypothesis.strategies as st # type: ignore[import-not-found] Review Comment: Agreed. TODO(rok). ########## python/pyarrow/tests/test_convert_builtin.py: ########## @@ -2452,6 +2457,7 @@ def test_nested_auto_chunking(ty, char): } [email protected] Review Comment: Opened an issue here https://github.com/apache/arrow/issues/55458. We can remove these, TODO(rok). ########## python/pyproject.toml: ########## Review Comment: Ah, we should remove excludes then, TODO(rok). ########## docs/source/developers/python/development.rst: ########## @@ -114,6 +114,18 @@ Running Type Checkers We support multiple type checkers. Their configurations are in ``pyproject.toml``. +The easiest way to run all of them is with ``pre-commit``, which installs +the same versions of the type checkers and stub packages for everyone. This +also lints and formats the stubs with ruff: + +.. code-block:: + + $ pre-commit run --hook-stage manual python-stubs --all-files Review Comment: Yeah, let's duplicate now, and reduce at the end of this effort if needed. TODO(rok) ########## ci/scripts/python_test_type_annotations.sh: ########## Review Comment: Let's pin here ad-hoc and simplify later. TODO(rok) ########## .pre-commit-config.yaml: ########## @@ -136,15 +136,101 @@ repos: - id: flake8 alias: python name: Python Lint + additional_dependencies: + - flake8-pyi Review Comment: Let's remove. TODO(rok). ########## python/pyarrow-stubs/pyarrow/__lib/array.pyi: ########## @@ -0,0 +1,613 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import sys + +from collections.abc import Iterable, Iterator, Sequence + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +from typing import Any, Literal + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] + +from pyarrow._compute import CastOptions +from pyarrow._stubs_typing import ( + ArrayLike, + Indices, + Mask, + MapsAsPydicts, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsPyArrowArray, +) +from pyarrow.lib import ( + Buffer, + Device, + MemoryManager, + MemoryPool, + Tensor, + _Weakrefable, +) +from typing_extensions import deprecated +import builtins + +from .scalar import Int64Scalar, Scalar +from .device import DeviceAllocationType +from .types import BaseExtensionType, DataType, Field, StructType +from pyarrow._stubs_typing import NullableCollection + +def array( + values: Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray, + type: Any | None = None, + mask: Mask | pd.Series[bool] | None = None, Review Comment: Removing `pd.Series` here seems like a good idea. I'd slightly prefer `Iterable` over custom protocol too. ########## python/pyarrow-stubs/pyarrow/__lib/types.pyi: ########## @@ -0,0 +1,700 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Any, Literal, TypeAlias, overload + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] +from typing_extensions import CapsuleType, Self, deprecated + +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow.lib import ( + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + Table, +) + +class _Weakrefable: ... +class _Metadata(_Weakrefable): ... + +class DataType(_Weakrefable): + def field(self, i: int) -> Field: ... + @property + def id(self) -> int: ... + @property + def bit_width(self) -> int: ... + @property + def byte_width(self) -> int: ... + @property + def num_fields(self) -> int: ... + @property + def num_buffers(self) -> int: ... + @property + def has_variadic_buffers(self) -> bool: ... + def __hash__(self) -> int: ... + def equals( + self, other: DataType | str, *, check_metadata: bool = False + ) -> bool: ... + def to_pandas_dtype(self) -> _PandasDtype: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> DataType: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> DataType: ... + +_PandasDtype: TypeAlias = ( + type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype +) + +class NullType(DataType): ... +class BoolType(DataType): ... +class UInt8Type(DataType): ... +class Int8Type(DataType): ... +class UInt16Type(DataType): ... +class Int16Type(DataType): ... +class UInt32Type(DataType): ... +class Int32Type(DataType): ... +class UInt64Type(DataType): ... +class Int64Type(DataType): ... +class Float16Type(DataType): ... +class Float32Type(DataType): ... +class Float64Type(DataType): ... +class Date32Type(DataType): ... +class Date64Type(DataType): ... +class MonthDayNanoIntervalType(DataType): ... +class StringType(DataType): ... +class LargeStringType(DataType): ... +class StringViewType(DataType): ... +class BinaryType(DataType): ... +class LargeBinaryType(DataType): ... +class BinaryViewType(DataType): ... + +_Time32Unit: TypeAlias = Literal["s", "ms"] +_Time64Unit: TypeAlias = Literal["us", "ns"] + +class TimestampType(DataType): + @property + def unit(self) -> TimeUnit: ... + @property + def tz(self) -> str | None: ... + +class Time32Type(DataType): + @property + def unit(self) -> _Time32Unit: ... + +class Time64Type(DataType): + @property + def unit(self) -> _Time64Unit: ... + +class DurationType(DataType): + @property + def unit(self) -> TimeUnit: ... + +class FixedSizeBinaryType(DataType): ... + +class Decimal32Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal64Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal128Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal256Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class ListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class ListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class FixedSizeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + @property + def list_size(self) -> int: ... + +class DictionaryMemo(_Weakrefable): ... + +_IndexType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type +) + +class DictionaryType(DataType): + @property + def ordered(self) -> bool: ... + @property + def index_type(self) -> _IndexType: ... + @property + def value_type(self) -> DataType: ... + +class MapType(DataType): + @property + def key_field(self) -> Field: ... + @property + def key_type(self) -> DataType: ... + @property + def item_field(self) -> Field: ... + @property + def item_type(self) -> DataType: ... + @property + def keys_sorted(self) -> bool: ... + +class StructType(DataType): + def get_field_index(self, name: str | bytes) -> int: ... + def field(self, i: int | str | bytes) -> Field: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + + __getitem__ = field + @property + def names(self) -> list[str]: ... + @property + def fields(self) -> list[Field]: ... + +class UnionType(DataType): + @property + def mode(self) -> Literal["sparse", "dense"]: ... + @property + def type_codes(self) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + def field(self, i: int) -> Field: ... + + __getitem__ = field + +class SparseUnionType(UnionType): + @property + def mode(self) -> Literal["sparse"]: ... + +class DenseUnionType(UnionType): + @property + def mode(self) -> Literal["dense"]: ... + +_RunEndType: TypeAlias = Int16Type | Int32Type | Int64Type + +class RunEndEncodedType(DataType): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> DataType: ... + +class BaseExtensionType(DataType): + def __arrow_ext_class__(self) -> type[ExtensionArray]: ... + def __arrow_ext_scalar_class__(self) -> type[ExtensionScalar]: ... + @property + def extension_name(self) -> str: ... + @property + def storage_type(self) -> DataType: ... + def wrap_array( + self, storage: Array | ChunkedArray + ) -> ExtensionArray | ChunkedArray: ... + +class ExtensionType(BaseExtensionType): + def __init__(self, storage_type: DataType, extension_name: str) -> None: ... + def __arrow_ext_serialize__(self) -> bytes: ... + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: DataType, serialized: bytes + ) -> Self: ... + +class FixedShapeTensorType(BaseExtensionType): + @property + def value_type(self) -> DataType: ... + @property + def shape(self) -> list[int]: ... + @property + def dim_names(self) -> list[str] | None: ... + @property + def permutation(self) -> list[int] | None: ... + +class Bool8Type(BaseExtensionType): ... +class UuidType(BaseExtensionType): ... +class JsonType(BaseExtensionType): ... + +class OpaqueType(BaseExtensionType): + @property + def type_name(self) -> str: ... + @property + def vendor_name(self) -> str: ... + +class UnknownExtensionType(ExtensionType): + def __init__(self, storage_type: DataType, serialized: bytes) -> None: ... + +def register_extension_type(ext_type: BaseExtensionType) -> None: ... +def unregister_extension_type(type_name: str) -> None: ... + +_FieldName: TypeAlias = str | bytes +_StrOrBytes: TypeAlias = str | bytes +_MetadataMapping: TypeAlias = ( + Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] +) +_MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] +_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable | None +_DataTypeAlias: TypeAlias = Literal[ + "null", + "bool", + "boolean", + "i1", + "int8", + "i2", + "int16", + "i4", + "int32", + "i8", + "int64", + "u1", + "uint8", + "u2", + "uint16", + "u4", + "uint32", + "u8", + "uint64", + "f2", + "halffloat", + "float16", + "f4", + "float", + "float32", + "f8", + "double", + "float64", + "string", + "str", + "utf8", + "binary", + "large_string", + "large_str", + "large_utf8", + "large_binary", + "binary_view", + "string_view", + "date32", + "date64", + "date32[day]", + "date64[ms]", + "time32[s]", + "time32[ms]", + "time64[us]", + "time64[ns]", + "timestamp[s]", + "timestamp[ms]", + "timestamp[us]", + "timestamp[ns]", + "duration[s]", + "duration[ms]", + "duration[us]", + "duration[ns]", + "month_day_nano_interval", +] +_DataTypeAliasInput: TypeAlias = _DataTypeAlias | str +_DataTypeLike: TypeAlias = DataType | _DataTypeAliasInput + +class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): + def __init__( + self, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, + **kwargs: str | bytes, + ) -> None: ... + def equals(self, other: KeyValueMetadata) -> bool: ... + def __len__(self) -> int: ... + def __contains__(self, /, __key: object) -> bool: ... + def __getitem__(self, /, __key: str | bytes) -> bytes: ... Review Comment: I'd try to get as close to what we have in Arrow. ########## python/pyarrow-stubs/pyarrow/__lib/array.pyi: ########## @@ -0,0 +1,606 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from collections.abc import Iterable, Iterator, Sequence + +from typing import Any, Literal, Self + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] + +from pyarrow._compute import CastOptions +from pyarrow._stubs_typing import ( + ArrayLike, + Indices, + Mask, + MapsAsPydicts, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsPyArrowArray, +) +from pyarrow.lib import ( + Buffer, + Device, + MemoryManager, + MemoryPool, + Tensor, + _Weakrefable, +) +from typing_extensions import deprecated +import builtins + +from .scalar import Int64Scalar, Scalar +from .device import DeviceAllocationType +from .types import BaseExtensionType, DataType, Field, StructType +from pyarrow._stubs_typing import NullableCollection + +def array( + values: Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray, + type: Any | None = None, + mask: Mask | pd.Series[bool] | None = None, + size: int | None = None, + from_pandas: bool | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... +def asarray( + values: Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray, + type: Any | None = None, +) -> ArrayLike: ... +def nulls( + size: int, + type: Any | None = None, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... +def repeat( + value: Any, + size: int, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... +def infer_type( + values: Iterable[Any], mask: Mask | None = None, from_pandas: bool = False +) -> DataType: ... + +class ArrayStatistics(_Weakrefable): + @property + def null_count(self) -> int | None: ... + @property + def distinct_count(self) -> int | None: ... + @property + def is_null_count_exact(self) -> bool | None: ... + @property + def is_distinct_count_exact(self) -> bool | None: ... + @property + def min(self) -> Any | None: ... + @property + def is_min_exact(self) -> bool | None: ... + @property + def max(self) -> Any | None: ... + @property + def is_max_exact(self) -> bool | None: ... + +class _PandasConvertible(_Weakrefable): + def to_pandas( + self, + memory_pool: MemoryPool | None = None, + categories: list | tuple | None = None, + strings_to_categorical: bool = False, + zero_copy_only: bool = False, + integer_object_nulls: bool = False, + date_as_object: bool = True, + timestamp_as_object: bool = False, + use_threads: bool = True, + deduplicate_objects: bool = True, + ignore_metadata: bool = False, + safe: bool = True, + split_blocks: bool = False, + self_destruct: bool = False, + maps_as_pydicts: MapsAsPydicts | None = None, + types_mapper: Any = None, # Callable[[DataType], ExtensionDtype | None] | None + coerce_temporal_nanoseconds: bool = False, + ) -> pd.Series | pd.DataFrame: ... + +class Array(_PandasConvertible): + def diff(self, other: Self) -> str: ... Review Comment: Should we be replacing `Self` with concrete types (`Array` here) or keep `Self`. We are not consistent right now and don't think we discussed this yet. @AlenkaF @tadeja @jorisvandenbossche ? ########## python/pyarrow-stubs/pyarrow/__lib/array.pyi: ########## @@ -0,0 +1,613 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import sys + +from collections.abc import Iterable, Iterator, Sequence + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +from typing import Any, Literal + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] + +from pyarrow._compute import CastOptions +from pyarrow._stubs_typing import ( + ArrayLike, + Indices, + Mask, + MapsAsPydicts, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsPyArrowArray, +) +from pyarrow.lib import ( + Buffer, + Device, + MemoryManager, + MemoryPool, + Tensor, + _Weakrefable, +) +from typing_extensions import deprecated +import builtins + +from .scalar import Int64Scalar, Scalar +from .device import DeviceAllocationType +from .types import BaseExtensionType, DataType, Field, StructType +from pyarrow._stubs_typing import NullableCollection + +def array( + values: Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray, + type: Any | None = None, + mask: Mask | pd.Series[bool] | None = None, + size: int | None = None, + from_pandas: bool | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... +def asarray( + values: Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray, + type: Any | None = None, +) -> ArrayLike: ... +def nulls( + size: int, + type: Any | None = None, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... +def repeat( + value: Any, + size: int, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... +def infer_type( + values: Iterable[Any], mask: Mask | None = None, from_pandas: bool = False +) -> DataType: ... + +class ArrayStatistics(_Weakrefable): + @property + def null_count(self) -> int | None: ... + @property + def distinct_count(self) -> int | None: ... + @property + def is_null_count_exact(self) -> bool | None: ... + @property + def is_distinct_count_exact(self) -> bool | None: ... + @property + def min(self) -> Any | None: ... + @property + def is_min_exact(self) -> bool | None: ... + @property + def max(self) -> Any | None: ... + @property + def is_max_exact(self) -> bool | None: ... + +class _PandasConvertible(_Weakrefable): + def to_pandas( + self, + memory_pool: MemoryPool | None = None, + categories: list | tuple | None = None, + strings_to_categorical: bool = False, + zero_copy_only: bool = False, + integer_object_nulls: bool = False, + date_as_object: bool = True, + timestamp_as_object: bool = False, + use_threads: bool = True, + deduplicate_objects: bool = True, + ignore_metadata: bool = False, + safe: bool = True, + split_blocks: bool = False, + self_destruct: bool = False, + maps_as_pydicts: MapsAsPydicts | None = None, + types_mapper: Any = None, # Callable[[DataType], ExtensionDtype | None] | None + coerce_temporal_nanoseconds: bool = False, + ) -> pd.Series | pd.DataFrame: ... + +class Array(_PandasConvertible): + def diff(self, other: Self) -> str: ... + def cast( + self, + target_type: DataType | str, + safe: bool = True, + options: CastOptions | None = None, + memory_pool: MemoryPool | None = None, + ) -> Array: ... + def view(self, target_type: DataType) -> Array: ... + def sum(self, **kwargs) -> Scalar: ... + @property + def type(self) -> DataType: ... + def unique(self) -> Self: ... + def dictionary_encode(self, null_encoding: str = "mask") -> DictionaryArray: ... + def value_counts(self) -> StructArray: ... + @staticmethod + def from_pandas( + obj: pd.Series | np.ndarray | ArrayLike, + mask: Mask | None = None, + type: DataType | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, + ) -> Array: ... + @staticmethod + def from_buffers( + type: DataType, + length: int, + buffers: Sequence[Buffer | None], Review Comment: Semantically these need to come in order and Collection is not necessarily ordered? ########## python/pyarrow-stubs/pyarrow/__lib/types.pyi: ########## @@ -0,0 +1,700 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Any, Literal, TypeAlias, overload + +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] +from typing_extensions import CapsuleType, Self, deprecated + +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow.lib import ( + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + Table, +) + +class _Weakrefable: ... +class _Metadata(_Weakrefable): ... + +class DataType(_Weakrefable): + def field(self, i: int) -> Field: ... + @property + def id(self) -> int: ... + @property + def bit_width(self) -> int: ... + @property + def byte_width(self) -> int: ... + @property + def num_fields(self) -> int: ... + @property + def num_buffers(self) -> int: ... + @property + def has_variadic_buffers(self) -> bool: ... + def __hash__(self) -> int: ... + def equals( + self, other: DataType | str, *, check_metadata: bool = False + ) -> bool: ... + def to_pandas_dtype(self) -> _PandasDtype: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> DataType: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> DataType: ... + +_PandasDtype: TypeAlias = ( + type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype +) + +class NullType(DataType): ... +class BoolType(DataType): ... +class UInt8Type(DataType): ... +class Int8Type(DataType): ... +class UInt16Type(DataType): ... +class Int16Type(DataType): ... +class UInt32Type(DataType): ... +class Int32Type(DataType): ... +class UInt64Type(DataType): ... +class Int64Type(DataType): ... +class Float16Type(DataType): ... +class Float32Type(DataType): ... +class Float64Type(DataType): ... +class Date32Type(DataType): ... +class Date64Type(DataType): ... +class MonthDayNanoIntervalType(DataType): ... +class StringType(DataType): ... +class LargeStringType(DataType): ... +class StringViewType(DataType): ... +class BinaryType(DataType): ... +class LargeBinaryType(DataType): ... +class BinaryViewType(DataType): ... + +_Time32Unit: TypeAlias = Literal["s", "ms"] +_Time64Unit: TypeAlias = Literal["us", "ns"] + +class TimestampType(DataType): + @property + def unit(self) -> TimeUnit: ... + @property + def tz(self) -> str | None: ... + +class Time32Type(DataType): + @property + def unit(self) -> _Time32Unit: ... + +class Time64Type(DataType): + @property + def unit(self) -> _Time64Unit: ... + +class DurationType(DataType): + @property + def unit(self) -> TimeUnit: ... + +class FixedSizeBinaryType(DataType): ... + +class Decimal32Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal64Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal128Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal256Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class ListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class ListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class FixedSizeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + @property + def list_size(self) -> int: ... + +class DictionaryMemo(_Weakrefable): ... + +_IndexType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type +) + +class DictionaryType(DataType): + @property + def ordered(self) -> bool: ... + @property + def index_type(self) -> _IndexType: ... + @property + def value_type(self) -> DataType: ... + +class MapType(DataType): + @property + def key_field(self) -> Field: ... + @property + def key_type(self) -> DataType: ... + @property + def item_field(self) -> Field: ... + @property + def item_type(self) -> DataType: ... + @property + def keys_sorted(self) -> bool: ... + +class StructType(DataType): + def get_field_index(self, name: str | bytes) -> int: ... + def field(self, i: int | str | bytes) -> Field: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + + __getitem__ = field + @property + def names(self) -> list[str]: ... + @property + def fields(self) -> list[Field]: ... + +class UnionType(DataType): + @property + def mode(self) -> Literal["sparse", "dense"]: ... + @property + def type_codes(self) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + def field(self, i: int) -> Field: ... + + __getitem__ = field + +class SparseUnionType(UnionType): + @property + def mode(self) -> Literal["sparse"]: ... + +class DenseUnionType(UnionType): + @property + def mode(self) -> Literal["dense"]: ... + +_RunEndType: TypeAlias = Int16Type | Int32Type | Int64Type + +class RunEndEncodedType(DataType): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> DataType: ... + +class BaseExtensionType(DataType): + def __arrow_ext_class__(self) -> type[ExtensionArray]: ... + def __arrow_ext_scalar_class__(self) -> type[ExtensionScalar]: ... + @property + def extension_name(self) -> str: ... + @property + def storage_type(self) -> DataType: ... + def wrap_array( + self, storage: Array | ChunkedArray + ) -> ExtensionArray | ChunkedArray: ... + +class ExtensionType(BaseExtensionType): + def __init__(self, storage_type: DataType, extension_name: str) -> None: ... + def __arrow_ext_serialize__(self) -> bytes: ... + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: DataType, serialized: bytes + ) -> Self: ... + +class FixedShapeTensorType(BaseExtensionType): + @property + def value_type(self) -> DataType: ... + @property + def shape(self) -> list[int]: ... + @property + def dim_names(self) -> list[str] | None: ... + @property + def permutation(self) -> list[int] | None: ... + +class Bool8Type(BaseExtensionType): ... +class UuidType(BaseExtensionType): ... +class JsonType(BaseExtensionType): ... + +class OpaqueType(BaseExtensionType): + @property + def type_name(self) -> str: ... + @property + def vendor_name(self) -> str: ... + +class UnknownExtensionType(ExtensionType): + def __init__(self, storage_type: DataType, serialized: bytes) -> None: ... + +def register_extension_type(ext_type: BaseExtensionType) -> None: ... +def unregister_extension_type(type_name: str) -> None: ... + +_FieldName: TypeAlias = str | bytes +_StrOrBytes: TypeAlias = str | bytes +_MetadataMapping: TypeAlias = ( + Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] +) +_MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] +_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable | None +_DataTypeAlias: TypeAlias = Literal[ + "null", + "bool", + "boolean", + "i1", + "int8", + "i2", + "int16", + "i4", + "int32", + "i8", + "int64", + "u1", + "uint8", + "u2", + "uint16", + "u4", + "uint32", + "u8", + "uint64", + "f2", + "halffloat", + "float16", + "f4", + "float", + "float32", + "f8", + "double", + "float64", + "string", + "str", + "utf8", + "binary", + "large_string", + "large_str", + "large_utf8", + "large_binary", + "binary_view", + "string_view", + "date32", + "date64", + "date32[day]", + "date64[ms]", + "time32[s]", + "time32[ms]", + "time64[us]", + "time64[ns]", + "timestamp[s]", + "timestamp[ms]", + "timestamp[us]", + "timestamp[ns]", + "duration[s]", + "duration[ms]", + "duration[us]", + "duration[ns]", + "month_day_nano_interval", +] +_DataTypeAliasInput: TypeAlias = _DataTypeAlias | str +_DataTypeLike: TypeAlias = DataType | _DataTypeAliasInput + +class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): + def __init__( + self, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, + **kwargs: str | bytes, + ) -> None: ... + def equals(self, other: KeyValueMetadata) -> bool: ... + def __len__(self) -> int: ... + def __contains__(self, /, __key: object) -> bool: ... + def __getitem__(self, /, __key: str | bytes) -> bytes: ... + def __iter__(self) -> Iterator[bytes]: ... + def key(self, i: int) -> bytes: ... + def value(self, i: int) -> bytes: ... + # The runtime yields generators instead of Mapping views + def keys(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def values(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def items(self) -> Iterator[tuple[bytes, bytes]]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def get_all(self, key: str | bytes) -> list[bytes]: ... + def to_dict(self) -> dict[bytes, bytes]: ... + +class Field(_Weakrefable): + def equals(self, other: Field, check_metadata: bool = False) -> bool: ... + def __hash__(self) -> int: ... + @property + def nullable(self) -> bool: ... + @property + def name(self) -> str: ... + @property + def metadata(self) -> dict[bytes, bytes] | None: ... + @property + def type(self) -> DataType: ... + def with_metadata( + self, + metadata: _KeyValueMetadataInput | KeyValueMetadata, + ) -> Field: ... + def remove_metadata(self) -> Field: ... + def with_type(self, new_type: DataType) -> Field: ... + def with_name(self, name: str | bytes) -> Field: ... + def with_nullable(self, nullable: bool) -> Field: ... + def flatten(self) -> list[Field]: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> Field: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Field: ... + +_FieldTuple: TypeAlias = ( + tuple[_FieldName, _DataTypeLike] + | tuple[_FieldName, _DataTypeLike, bool | None] + | tuple[ + _FieldName, + _DataTypeLike, + bool | None, + _KeyValueMetadataInput | KeyValueMetadata, + ] +) +_StructFieldsInput: TypeAlias = ( + Iterable[Field | _FieldTuple] | Mapping[_FieldName, _DataTypeLike] +) +_SchemaFieldsInput: TypeAlias = _StructFieldsInput + +class Schema(_Weakrefable): + def __len__(self) -> int: ... + def __getitem__(self, key: int) -> Field: ... + + _field = __getitem__ + def __iter__(self) -> Iterator[Field]: ... + def __hash__(self) -> int: ... + def __sizeof__(self) -> int: ... + @property + def pandas_metadata(self) -> dict[str, Any] | None: ... + @property + def names(self) -> list[str]: ... + @property + def types(self) -> list[DataType]: ... + @property + def metadata(self) -> dict[bytes, bytes] | None: ... + def empty_table(self) -> Table: ... + def equals(self, other: Schema, check_metadata: bool = False) -> bool: ... + @classmethod + def from_pandas( + cls, df: pd.DataFrame, preserve_index: bool | None = None + ) -> Schema: ... + def field(self, i: int | str | bytes) -> Field: ... + @deprecated("Use 'field' instead") + def field_by_name(self, name: str | bytes) -> Field | None: ... + def get_field_index(self, name: str | bytes) -> int: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def append(self, field: Field) -> Schema: ... + def insert(self, i: int, field: Field) -> Schema: ... + def remove(self, i: int) -> Schema: ... + def set(self, i: int, field: Field) -> Schema: ... + @deprecated("Use 'with_metadata' instead") + def add_metadata( + self, metadata: _KeyValueMetadataInput | KeyValueMetadata + ) -> Schema: ... + def with_metadata( + self, metadata: _KeyValueMetadataInput | KeyValueMetadata + ) -> Schema: ... + def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... + def remove_metadata(self) -> Schema: ... + def to_string( + self, + truncate_metadata: bool = True, + show_field_metadata: bool = True, + show_schema_metadata: bool = True, + element_size_limit: int = 100, + ) -> str: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> Schema: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Schema: ... + +def unify_schemas( + schemas: Iterable[Schema], + *, + promote_options: Literal["default", "permissive"] = "default", +) -> Schema: ... +@overload +def field( + name: SupportsArrowSchema, + type: None = None, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata = None, +) -> Field: ... +@overload +def field( + name: _FieldName, + type: _DataTypeLike, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata = None, +) -> Field: ... +def null() -> NullType: ... +def bool_() -> BoolType: ... +def uint8() -> UInt8Type: ... +def int8() -> Int8Type: ... +def uint16() -> UInt16Type: ... +def int16() -> Int16Type: ... +def uint32() -> UInt32Type: ... +def int32() -> Int32Type: ... +def int64() -> Int64Type: ... +def uint64() -> UInt64Type: ... +def timestamp(unit: TimeUnit | str, tz: str | None = None) -> TimestampType: ... Review Comment: If we can't solve this in a nice way, I'd relax with strings to not inconvenience users. It's not great, but casts to silence typecheckers are not great downstream. ########## python/pyarrow/tests/test_array.py: ########## @@ -18,20 +18,24 @@ from collections.abc import Iterable import datetime import decimal -import hypothesis as h -import hypothesis.strategies as st +import hypothesis as h # type: ignore[import-not-found] +import hypothesis.strategies as st # type: ignore[import-not-found] import itertools import operator -import pytest +import pytest # type: ignore[import-not-found] import struct import subprocess import sys import weakref +from typing import TYPE_CHECKING -try: +if TYPE_CHECKING: import numpy as np -except ImportError: - np = None +else: + try: + import numpy as np + except ImportError: + np = None Review Comment: Will check. TODO(rok) ########## python/pyarrow/tests/test_array.py: ########## @@ -2868,7 +2881,7 @@ def test_interval_array_from_relativedelta(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([13, 8, + pa.MonthDayNano([13, 8, # type: ignore[arg-type] Review Comment: `MonthDayNano` is not annotated in this PR, hence is annotated `Any`, https://github.com/apache/arrow/pull/55264/changes#r4515157324 to add it. ########## python/pyarrow/tests/test_array.py: ########## @@ -1296,24 +1303,24 @@ def test_map_from_arrays(): keys = pa.array(pykeys, type='binary') items = pa.array(pyitems, type='i4') - result = pa.MapArray.from_arrays(offsets, keys, items) + result = pa.MapArray.from_arrays(offsets, keys, items) # type: ignore[arg-type] expected = pa.array(pyentries, type=pa.map_(pa.binary(), pa.int32())) assert result.equals(expected) # pass in the type explicitly - result = pa.MapArray.from_arrays(offsets, keys, items, pa.map_( - keys.type, - items.type - )) + result = pa.MapArray.from_arrays(offsets, keys, items, # type: ignore[arg-type] Review Comment: I think we'd get `too many positional arguments`, but let's remove it and see. TODO(rok). ########## python/pyproject.toml: ########## Review Comment: I don't know how to answer right now, we need to test this thoroughly. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
