jorisvandenbossche commented on code in PR #55264: URL: https://github.com/apache/arrow/pull/55264#discussion_r4481425759
########## 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 don't know if this is hallucinated or because someone of you actually run into an issue with this, but using `python -m pytest` instead of `pytest` with the example above works just fine for me. (in general this is just about running pytest, so that should generally work the same as running the normal pyarrow unit tests? And there we don't say anything like this: https://arrow.apache.org/docs/developers/python/development.html#unit-testing) ########## 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: Is there a reason this is all written with `lib` instead of the public import path? (to be honest, if using `pa.` for that, this is only about a single character being saved ;), but in general in tests we want to use public imports; that also makes things easier to copy or paste from interactive testing) ########## python/pyarrow-stubs/tests/_check.py: ########## @@ -0,0 +1,58 @@ +# 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. + +"""Helpers for the type annotation tests in this directory. + +Each test pairs a static and a runtime check of the same expression:: + + check(assert_type(pa.lib.list_(pa.lib.int32()), pa.lib.ListType), + pa.lib.ListType) + +``assert_type`` is verified by every type checker we run in CI (mypy, +pyright and ty) and is a no-op at runtime. ``check`` runs under pytest and +verifies that the value really is an instance of the given class, so the +stubs can't drift away from what pyarrow actually returns. + +Code that type checkers should *reject* goes under +``if TYPE_CHECKING_INVALID_USAGE:`` with a ``# type: ignore[<mypy code>]`` +comment. The block never runs, and mypy and pyright are configured to +report unused ignores in these tests, so the test fails if the usage stops +being an error. +""" + +import sys +from typing import Any, Final, TYPE_CHECKING, TypeVar + +if sys.version_info >= (3, 11): + from typing import assert_type +else: + from typing_extensions import assert_type Review Comment: pyarrow is already on `>=3.11`, so this can be simplified? ########## 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: It feels strange that we have an import here that looks like valid code but that _only_ works at type checking time? (I would expect the stubs to be a 1:1 mapping to the actual code, and additional (helper) types to be defined elsewhere) If those classes are not actually defined in `pyarrow.lib`, should we define them in a `pyarrow.typing`-like module instead? ########## 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 + == types.TypesEnum.INTERVAL_MONTH_DAY_NANO) + + +def test_data_type_hashable() -> None: + data_type: Hashable = lib.int32() Review Comment: ```suggestion data_type = lib.int32() ``` it should not be needed to explicitly declare this as Hashable? (I would expect the static typing to infer 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 + == types.TypesEnum.INTERVAL_MONTH_DAY_NANO) + + +def test_data_type_hashable() -> None: + data_type: Hashable = lib.int32() + check(assert_type(hash(data_type), int), int) + + +def test_field() -> None: + field = lib.field("a", lib.int32()) + check(assert_type(field, lib.Field), lib.Field) + check(assert_type(field.name, str), str) + check(assert_type(field.type, lib.DataType), lib.DataType) + check(assert_type(field.nullable, bool), bool) + check(assert_type(field.with_nullable(False), lib.Field), lib.Field) + check(assert_type(lib.field("b", "string"), lib.Field), lib.Field) + + +def test_schema() -> None: + schema = lib.schema([("a", lib.int32()), lib.field("b", lib.string())]) + check(assert_type(schema, lib.Schema), lib.Schema) + check(assert_type(schema.names, list[str]), list) + check(assert_type(schema.types, list[lib.DataType]), list) + check(assert_type(schema.field("a"), lib.Field), lib.Field) + check(assert_type(schema[0], lib.Field), lib.Field) + + +def test_key_value_metadata() -> None: + # The runtime yields generators instead of Mapping views (duplicated keys) + metadata = lib.KeyValueMetadata([(b"a", b"1"), (b"a", b"2")]) + check(assert_type(metadata.keys(), Iterator[bytes]), Iterator) + check(assert_type(metadata.values(), Iterator[bytes]), Iterator) + check(assert_type(metadata.items(), Iterator[tuple[bytes, bytes]]), + Iterator) + + +def test_type_predicates_narrow() -> None: + lst = _erase(lib.list_(lib.int32())) + assert types.is_list(lst) + check(assert_type(lst, lib.ListType), lib.ListType) + + ts = _erase(lib.timestamp("s")) + assert types.is_timestamp(ts) + check(assert_type(ts, lib.TimestampType), lib.TimestampType) + + dec = _erase(lib.decimal64(5)) + assert types.is_decimal(dec) + check( + assert_type( + dec, + lib.Decimal32Type | lib.Decimal64Type | lib.Decimal128Type + | lib.Decimal256Type, + ), + lib.Decimal64Type, + ) + + # Decimal types subclass FixedSizeBinaryType, but don't match it, so + # they must not be narrowed away when the check fails. + dec128 = _erase(lib.decimal128(5, 2)) + if not types.is_fixed_size_binary(dec128): + assert types.is_decimal(dec128) + check( + assert_type( + dec128, + lib.Decimal32Type | lib.Decimal64Type | lib.Decimal128Type + | lib.Decimal256Type, + ), + lib.Decimal128Type, + ) + + # Run-end encoded types are nested, so this doesn't narrow to Never. + ree = lib.run_end_encoded(lib.int32(), lib.string()) Review Comment: The `_erase` is missing here? ########## 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: Instead of adding those ignores everywhere, can we just ensure we install those required test dependencies in the env where we do type checking? ########## 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: Why is this failing? (this is not testing a failure case, as the ignores above) ########## python/pyproject.toml: ########## @@ -149,3 +174,32 @@ exclude = [ "examples", "scripts", ] + +# numpy, pandas and scipy are optional dependencies. A `ty: ignore` comment on +# their imports would be reported as unused whenever they are installed. +[[tool.ty.overrides]] +include = ["pyarrow-stubs/pyarrow/**"] +[tool.ty.overrides.rules] +unresolved-import = "ignore" + +# ty doesn't honor `# type: ignore[<mypy code>]` comments, so disable the +# rules that those comments suppress for mypy. +[[tool.ty.overrides]] +include = ["pyarrow-stubs/pyarrow/__lib/array.pyi"] +[tool.ty.overrides.rules] +# Some array classes override methods with a different signature at runtime. +invalid-method-override = "ignore" + +[[tool.ty.overrides]] +include = ["pyarrow-stubs/tests/**"] +[tool.ty.overrides.rules] +# mypy and pyright check that the invalid usage in the tests is rejected. +invalid-argument-type = "ignore" + +# Lints and formats the type stubs (*.pyi); see .pre-commit-config.yaml. +[tool.ruff] +line-length = 88 Review Comment: Should this specifically be limited to the pyarrow-stubs directory? ########## 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: This currently fails for me with: ``` $ pre-commit run --hook-stage manual python-stubs --all-files Python Stubs Lint (ruff).................................................Passed Python Stubs Format (ruff)...............................................Passed Python Stubs Type Check (mypy)...........................................Failed - hook id: mypy - exit code: 1 pyarrow-stubs/pyarrow/__lib/table.pyi:57: error: Module "pyarrow.compute" has no attribute "Expression" [attr-defined] Found 1 error in 1 file (checked 22 source files) Python Stubs Type Check (pyright)........................................Failed - hook id: pyright - exit code: 1 /home/joris/scipy/repos/arrow/python/pyarrow-stubs/pyarrow/__lib/table.pyi /home/joris/scipy/repos/arrow/python/pyarrow-stubs/pyarrow/__lib/table.pyi:57:29 - error: "Expression" is unknown import symbol (reportAttributeAccessIssue) 1 error, 0 warnings, 0 informations Python Stubs Type Check (ty).............................................Passed ``` ########## 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: Why is `pd.Series` being listed here specifically? Also, in practice this does not seem to be accepted at runtime ########## ci/scripts/python_test_type_annotations.sh: ########## Review Comment: My experience with the pandas type checking is that this will quickly break, when not pinning the dependencies. At the moment those _are_ pinned in the pre-commit config, but so either we could reuse that instead of this script, or it might be good to define this environment somewhere else such that it can be reused? (or for now just duplicate the pinning) ########## 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) + + check(assert_type(arr.type, lib.DataType), lib.DataType) + check(assert_type(arr.null_count, int), int) + check(assert_type(len(arr), int), int) + check(assert_type(arr.to_pylist(), list[Any]), list) + # Scalars aren't annotated yet, so indexing returns Any | Int64Array. + check(assert_type(arr[0], Any | lib.Int64Array), lib.Int64Scalar) + check(assert_type(arr[1:], Any | lib.Int64Array), lib.Int64Array) + + # Methods returning the same kind of array keep the subclass. + check(assert_type(arr.slice(1), lib.Int64Array), lib.Int64Array) + check(assert_type(arr.fill_null(0), lib.Int64Array), lib.Int64Array) + check(assert_type(arr.drop_null(), lib.Int64Array), lib.Int64Array) + + check(assert_type(arr.is_null(), lib.BooleanArray), lib.BooleanArray) + check(assert_type(arr.cast(lib.float64()), lib.Array), lib.DoubleArray) + check(assert_type(arr.cast("float64"), lib.Array), lib.DoubleArray) + check(assert_type(lib.concat_arrays([arr, arr]), lib.Array), + lib.Int64Array) + check(assert_type(lib.Array.from_pandas([1, 2], None, lib.int64()), + lib.Array), + lib.Int64Array) + + other = check(cast(lib.StringArray, lib.array(["a", "b", None])), lib.StringArray) + check(assert_type(arr.equals(arr), bool), bool) + check(assert_type(arr.equals(other), bool), bool) + + +def test_array_dlpack() -> None: + arr = cast(lib.Int64Array, lib.array([1, 2])) + capsule = arr.__dlpack__( + max_version=(1, 0), + dl_device=arr.__dlpack_device__(), + copy=False, + ) + check(assert_type(capsule, Any), CapsuleType) + + +def test_import_from_c_capsule() -> None: + capsules = lib.array([1, 2]).__arrow_c_array__() + imported = lib.BooleanArray._import_from_c_capsule(*capsules) Review Comment: Is the BooleanArray (which contrasts with the actual type) intentional? In general I think the usage pattern for calling `_import_from_c_capsule` it to always do that using the base class `pa.Array` ########## python/pyarrow-stubs/tests/test_table.py: ########## @@ -0,0 +1,152 @@ +# 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 types import MappingProxyType + +import pytest + +import pyarrow.lib as lib + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + + +def test_table() -> None: + table = lib.table({"a": [1, 2], "b": ["x", "y"]}) + check(assert_type(table, lib.Table), lib.Table) + check(assert_type(table.num_rows, int), int) + check(assert_type(table.num_columns, int), int) + check(assert_type(table.shape, tuple[int, int]), tuple) + check(assert_type(table.column_names, list[str]), list) + check(assert_type(table.schema, lib.Schema), lib.Schema) + check(assert_type(table.field("a"), lib.Field), lib.Field) + + # Columns aren't typed per tabular class, so Table and RecordBatch both + # return the union here. + check(assert_type(table.column("a"), lib.Array | lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(table["a"], lib.Array | lib.ChunkedArray | lib.Table), Review Comment: The fact that this returns a union including `Table` is a TODO I suppose? (as that makes this quite useless) ########## python/pyarrow-stubs/tests/test_table.py: ########## @@ -0,0 +1,152 @@ +# 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 types import MappingProxyType + +import pytest + +import pyarrow.lib as lib + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + + +def test_table() -> None: + table = lib.table({"a": [1, 2], "b": ["x", "y"]}) + check(assert_type(table, lib.Table), lib.Table) + check(assert_type(table.num_rows, int), int) + check(assert_type(table.num_columns, int), int) + check(assert_type(table.shape, tuple[int, int]), tuple) + check(assert_type(table.column_names, list[str]), list) + check(assert_type(table.schema, lib.Schema), lib.Schema) + check(assert_type(table.field("a"), lib.Field), lib.Field) + + # Columns aren't typed per tabular class, so Table and RecordBatch both + # return the union here. + check(assert_type(table.column("a"), lib.Array | lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(table["a"], lib.Array | lib.ChunkedArray | lib.Table), + lib.ChunkedArray) + check(assert_type(table[0], lib.Array | lib.ChunkedArray | lib.Table), + lib.ChunkedArray) + check(assert_type(table[:1], lib.Array | lib.ChunkedArray | lib.Table), + lib.Table) + + check(assert_type(table.slice(0, 1), lib.Table), lib.Table) + check(assert_type(table.select(["a"]), lib.Table), lib.Table) + check(assert_type(table.select(["a", 1]), lib.Table), lib.Table) + check(assert_type(table.equals(None), bool), bool) + check(assert_type(table.combine_chunks(), lib.Table), lib.Table) + check(assert_type(table.to_batches(), list[lib.RecordBatch]), list) + check(assert_type(table.to_pydict(), dict[str, list]), dict) + check(assert_type(lib.Table.from_arrays([(1, 2)], names=["a"]), lib.Table), + lib.Table) + + chunked = lib.chunked_array([[1, 2]]) + check(assert_type(table.add_column(0, "c", chunked), lib.Table), lib.Table) + check(assert_type(table.add_column(0, "c", lib.array([1, 2])), lib.Table), + lib.Table) + check(assert_type(table.add_column(0, "c", [[1], [2]]), lib.Table), + lib.Table) + check(assert_type(table.add_column(0, "c", [lib.array([1, 2])]), lib.Table), + lib.Table) + + table = lib.table(([1],), names=["a"]) + check(assert_type(table, lib.Table), lib.Table) + + chunked_struct = lib.chunked_array([[{"x": 1}]]) + check(assert_type(lib.Table.from_struct_array(chunked_struct), lib.Table), + lib.Table) + + # The reader isn't annotated yet, but must not reject valid usage. + check(table.to_reader().read_all(), lib.Table) + for batch in table.to_reader(): + check(batch, lib.RecordBatch) + + +def test_record_batch() -> None: + batch = lib.record_batch({"a": [1, 2]}) # from a dict + check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + check(assert_type(batch.num_rows, int), int) + check(assert_type(batch.column(0), lib.Array | lib.ChunkedArray), + lib.Array) + check(assert_type(batch.slice(1), lib.RecordBatch), lib.RecordBatch) + check(assert_type(lib.Table.from_batches([batch]), lib.Table), lib.Table) + check(assert_type(lib.concat_batches([batch, batch]), lib.RecordBatch), + lib.RecordBatch) + check(assert_type(batch.add_column(0, "b", lib.array([3, 4])), + lib.RecordBatch), + lib.RecordBatch) + check(assert_type(batch.add_column(0, "b", [3, 4]), lib.RecordBatch), + lib.RecordBatch) + check(assert_type(batch.equals(batch), bool), bool) + # Any object can be compared, a non-batch returns False + check(assert_type(batch.equals(42), bool), bool) + + batch = lib.record_batch([lib.array([1])], names=["a"]) # from a list + check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + + struct_array = batch.to_struct_array() + check(assert_type(lib.RecordBatch.from_struct_array(struct_array), Review Comment: Also assert the type of `struct_array` here? ########## ci/scripts/python_test_type_annotations.sh: ########## @@ -36,3 +36,7 @@ cd "${pyarrow_dir}" mypy pyright --stats ty check --verbose --output-format concise + +# The type annotation tests are type checked above; also run them to verify +# the annotated types against the installed pyarrow. Review Comment: ```suggestion # the expected annotated types against the runtime types. ``` to make it clear that this pytest invocation does not actually check annotations (it only does because we repeat it twice, but so in itself just checks the runtime isinstance). (I was confused about what this pytest would do, in addition to running the type checkers) ########## 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: What is the value of those asserts for type checking? ########## .pre-commit-config.yaml: ########## @@ -136,15 +136,101 @@ repos: - id: flake8 alias: python name: Python Lint + additional_dependencies: + - flake8-pyi Review Comment: Is it needed to run flake8 on the pyi files if we already run ruff linting+formatting? (does flake8 at that point add any value?) ########## 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] Review Comment: This should not be needed? Where we run the type checks (pre-commit, the script), numpy is installed ########## 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: With an overload (for a later todo) this should be possible? ########## 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], + null_count: int = -1, + offset=0, + children: NullableCollection | None = None, + ) -> Array: ... + @property + def null_count(self) -> int: ... + @property + def nbytes(self) -> int: ... + def get_total_buffer_size(self) -> int: ... + def __sizeof__(self) -> int: ... + def __iter__(self) -> Iterator[Scalar]: ... + def to_string( + self, + *, + indent: int = 2, + top_level_indent: int = 0, + window: int = 10, + container_window: int = 2, + skip_new_lines: bool = False, + ) -> str: ... + + format = to_string + def equals(self, other: Array) -> bool: ... + def __len__(self) -> int: ... + def is_null(self, *, nan_is_null: bool = False) -> BooleanArray: ... + def is_nan(self) -> BooleanArray: ... + def is_valid(self) -> BooleanArray: ... + def fill_null(self, fill_value: Any) -> Self: ... + def __getitem__(self, key: int | builtins.slice) -> Scalar | Self: ... + def slice(self, offset: int = 0, length: int | None = None) -> Self: ... + def take(self, indices: Indices) -> Self: ... + def drop_null(self) -> Self: ... + def filter( + self, + mask: Mask, + *, + null_selection_behavior: Literal["drop", "emit_null"] = "drop", + ) -> Self: ... + def index( + self, + value: Any, + start: int | None = None, + end: int | None = None, + *, + memory_pool: MemoryPool | None = None, + ) -> Int64Scalar: ... + def sort(self, order: Order = "ascending", **kwargs) -> Self: ... + def __array__( + self, dtype: np.dtype | None = None, copy: bool | None = None + ) -> np.ndarray: ... + def to_numpy( + self, zero_copy_only: bool = True, writable: bool = False + ) -> np.ndarray: ... + def to_pylist( + self, + *, + maps_as_pydicts: Literal["lossy", "strict"] | None = None, + ) -> list[Any]: ... + + tolist = to_pylist + def validate(self, *, full: bool = False) -> None: ... + @property + def offset(self) -> int: ... + def buffers(self) -> list[Buffer | None]: ... + def copy_to(self, destination: MemoryManager | Device) -> Self: ... + def _export_to_c(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int, type: int | DataType) -> Array: ... + def __arrow_c_array__(self, requested_schema=None) -> Any: ... + @staticmethod + def _import_from_c_capsule(schema_capsule, array_capsule) -> Array: ... + def _export_to_c_device(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + @staticmethod + def _import_from_c_device(in_ptr: int, type: DataType | int) -> Array: ... + def __arrow_c_device_array__(self, requested_schema=None, **kwargs) -> Any: ... + @staticmethod + def _import_from_c_device_capsule(schema_capsule, array_capsule) -> Array: ... + def __dlpack__( + self, + *, + stream: int | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> Any: ... + def __dlpack_device__(self) -> tuple[int, int]: ... + @property + def device_type(self) -> DeviceAllocationType: ... + @property + def is_cpu(self) -> bool: ... + @property + def statistics(self) -> ArrayStatistics | None: ... + +class NullArray(Array): ... + +class BooleanArray(Array): + @property + def false_count(self) -> int: ... + @property + def true_count(self) -> int: ... + +class NumericArray(Array): ... +class IntegerArray(NumericArray): ... +class FloatingPointArray(NumericArray): ... +class Int8Array(IntegerArray): ... +class UInt8Array(IntegerArray): ... +class Int16Array(IntegerArray): ... +class UInt16Array(IntegerArray): ... +class Int32Array(IntegerArray): ... +class UInt32Array(IntegerArray): ... +class Int64Array(IntegerArray): ... +class UInt64Array(IntegerArray): ... +class Date32Array(NumericArray): ... +class Date64Array(NumericArray): ... +class TimestampArray(NumericArray): ... +class Time32Array(NumericArray): ... +class Time64Array(NumericArray): ... +class DurationArray(NumericArray): ... +class MonthDayNanoIntervalArray(Array): ... +class HalfFloatArray(FloatingPointArray): ... +class FloatArray(FloatingPointArray): ... +class DoubleArray(FloatingPointArray): ... +class FixedSizeBinaryArray(Array): ... +class Decimal32Array(FixedSizeBinaryArray): ... +class Decimal64Array(FixedSizeBinaryArray): ... +class Decimal128Array(FixedSizeBinaryArray): ... +class Decimal256Array(FixedSizeBinaryArray): ... + +class BaseListArray(Array): + def flatten(self, recursive: bool = False) -> Array: ... + def value_parent_indices(self) -> Int64Array: ... + def value_lengths(self) -> Int32Array | Int64Array: ... + +class ListArray(BaseListArray): + @classmethod + def from_arrays( + cls, + offsets: Int32Array | list[int] | list[int | None], + values: Array | list[Any], + *, + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> ListArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int32Array: ... + def value_lengths(self) -> Int32Array: ... + +class LargeListArray(BaseListArray): + @classmethod + def from_arrays( + cls, + offsets: Int64Array | list[int] | list[int | None], + values: Array, + *, + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> LargeListArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int64Array: ... + def value_lengths(self) -> Int64Array: ... + +class ListViewArray(BaseListArray): + @classmethod + def from_arrays( + cls, + offsets: Int32Array | list[int] | list[int | None], + sizes: Int32Array | list[int] | list[int | None], + values: Array, + *, + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> ListViewArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int32Array: ... + @property + def sizes(self) -> Int32Array: ... + def value_lengths(self) -> Int32Array: ... + +class LargeListViewArray(BaseListArray): + @classmethod + def from_arrays( + cls, + offsets: Int64Array | list[int] | list[int | None], + sizes: Int64Array | list[int] | list[int | None], + values: Array, + *, + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> LargeListViewArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int64Array: ... + @property + def sizes(self) -> Int64Array: ... + def value_lengths(self) -> Int64Array: ... + +class FixedSizeListArray(BaseListArray): + @classmethod + def from_arrays( + cls, + values: Array, + list_size: int | None = None, + *, + type: DataType | None = None, + mask: Mask | None = None, + ) -> FixedSizeListArray: ... + @property + def values(self) -> Array: ... + def value_lengths(self) -> Int32Array: ... + +class MapArray(ListArray): + @classmethod + def from_arrays( # type: ignore[override] + cls, + offsets: Int32Array | list[int] | list[int | None], + keys: Array | np.ndarray | list, + items: Array | np.ndarray | list, + *, + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> MapArray: ... + @property + def keys(self) -> Array: ... + @property + def items(self) -> Array: ... + +class UnionArray(Array): + @deprecated("Use fields() instead") + def child(self, pos: int) -> Field: ... + def field(self, pos: int) -> Array: ... + @property + def type_codes(self) -> Int8Array: ... + @property + def offsets(self) -> Int32Array: ... + @staticmethod + def from_dense( + types: Int8Array, + value_offsets: Int32Array, + children: NullableCollection, + field_names: list[str] | None = None, + type_codes: Int8Array | list[int] | None = None, + ) -> UnionArray: ... + @staticmethod + def from_sparse( + types: Int8Array, + children: NullableCollection, + field_names: list[str] | None = None, + type_codes: Int8Array | list[int] | None = None, + ) -> UnionArray: ... + +class StringArray(Array): + @staticmethod + def from_buffers( # type: ignore[override] + length: int, + value_offsets: Buffer, + data: Buffer, + null_bitmap: Buffer | None = None, + null_count: int | None = -1, + offset: int | None = 0, + ) -> StringArray: ... + +class LargeStringArray(Array): + @staticmethod + def from_buffers( # type: ignore[override] + length: int, + value_offsets: Buffer, + data: Buffer, + null_bitmap: Buffer | None = None, + null_count: int | None = -1, + offset: int | None = 0, + ) -> LargeStringArray: ... + +class StringViewArray(Array): ... + +class BinaryArray(Array): + @property + def total_values_length(self) -> int: ... + +class LargeBinaryArray(Array): + @property + def total_values_length(self) -> int: ... + +class BinaryViewArray(Array): ... + +class DictionaryArray(Array): + def dictionary_encode(self) -> Self: ... # type: ignore[override] + def dictionary_decode(self) -> Array: ... + @property + def indices(self) -> IntegerArray: ... + @property + def dictionary(self) -> Array: ... + @staticmethod + def from_buffers( # type: ignore[override] + type: DataType, + length: int, + buffers: list[Buffer], + dictionary: Array | np.ndarray | pd.Series, + null_count: int = -1, + offset: int = 0, + ) -> DictionaryArray: ... + @staticmethod + def from_arrays( + indices: Indices | Sequence[int | None], Review Comment: ```suggestion indices: Indices, ``` The Sequence is already included in the Indices alias? ########## 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, Review Comment: ```suggestion type: DataType | str | None = None, ``` ########## python/pyarrow-stubs/pyarrow/_stubs_typing.pyi: ########## @@ -0,0 +1,116 @@ +# 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 datetime as dt + +from _typeshed import ( + SupportsDunderGE, + SupportsDunderGT, + SupportsDunderLE, + SupportsDunderLT, +) +from collections.abc import Collection, Container, Iterator, Sequence, Sized +from decimal import Decimal +from typing import Any, Literal, Protocol, TypeAlias + +import numpy as np + +from numpy.typing import NDArray + +from pyarrow import lib +from pyarrow.lib import ChunkedArray + +ArrayLike: TypeAlias = Any +ScalarLike: TypeAlias = Any +Order: TypeAlias = Literal["ascending", "descending"] +JoinType: TypeAlias = Literal[ + "left semi", + "right semi", + "left anti", + "right anti", + "inner", + "left outer", + "right outer", + "full outer", +] +Compression: TypeAlias = Literal[ + "gzip", "bz2", "brotli", "lz4", "lz4_frame", "lz4_raw", "zstd", "snappy" +] +MapsAsPydicts: TypeAlias = Literal["lossy", "strict"] +NullEncoding: TypeAlias = Literal["mask", "encode"] +NullPlacement: TypeAlias = Literal["at_start", "at_end"] +NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] +TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] + +PyScalar: TypeAlias = ( + bool + | int + | float + | Decimal + | str + | bytes + | dt.date + | dt.datetime + | dt.time + | dt.timedelta +) +NumpyScalar: TypeAlias = np.generic[Any] + +Mask: TypeAlias = ( + Sequence[bool | None] | NDArray[np.bool_] | lib.BooleanArray | ChunkedArray +) +Indices: TypeAlias = ( + Sequence[int | None] | NDArray[np.integer[Any]] | lib.IntegerArray | ChunkedArray Review Comment: Similarly here, I think this should be `Array` instead of `IntegerArray`. Often your pyarrow array will only be inferred as Array, and you don't want to get constantly errors about that (or having to add type casts to avoid it) ########## 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: A nitpick, but I would find it cleaner to clearly separate `TYPE_CHECKING` imports from runtime imports, i.e. put the `if TYPE_CHECKING: ..` at the end after all normal imports. That means the numpy import is repeated (and an ignore is needed on this `None`), but makes does not have you reason about runtime imports being in the `else` clause ########## 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 + == types.TypesEnum.INTERVAL_MONTH_DAY_NANO) + + +def test_data_type_hashable() -> None: + data_type: Hashable = lib.int32() + check(assert_type(hash(data_type), int), int) + + +def test_field() -> None: + field = lib.field("a", lib.int32()) + check(assert_type(field, lib.Field), lib.Field) + check(assert_type(field.name, str), str) + check(assert_type(field.type, lib.DataType), lib.DataType) + check(assert_type(field.nullable, bool), bool) + check(assert_type(field.with_nullable(False), lib.Field), lib.Field) + check(assert_type(lib.field("b", "string"), lib.Field), lib.Field) + Review Comment: Add a test for `.metadata`? ########## 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: The fact that there is still the Int64Value here seems to indicate there is something wrong with this test? (i.e. never run? as this should just fail) I suppose the double mark does not actually work? (might be worth a separate issue) ########## 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: Why is this type ignore needed? (if I remove it, mypy still runs fine for me) ########## 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], + null_count: int = -1, + offset=0, + children: NullableCollection | None = None, Review Comment: ```suggestion children: Collection | None = None, ``` ? ########## python/pyarrow-stubs/pyarrow/__lib/table.pyi: ########## @@ -0,0 +1,600 @@ +# 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 + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self +if sys.version_info >= (3, 10): + from typing import TypeAlias +else: + from typing_extensions import 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, +) +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 = ( + Iterable[Array | Iterable[Any] | SupportsArrowStream | SupportsArrowArray] + | Array + | SupportsArrowArray + | SupportsArrowStream +) + +UnarySelector: TypeAlias = str +NullarySelector: TypeAlias = tuple[()] +NarySelector: TypeAlias = list[str] | tuple[str, ...] +ColumnSelector: TypeAlias = UnarySelector | NullarySelector | NarySelector + +class ChunkedArray(_PandasConvertible): + @property + def data(self) -> Self: ... Review Comment: Sidenote, we should probably just remove this attribute, it has been deprecated for years ########## 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: Strictly speaking, this does not need to be a `Sequence` (which has getitem), but `Collection` (or even `Sized`) is enough (for children below, which has the same requirements, a Collection is used) ########## 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: And actually the `if TYPE_CHECKING: import numpy as np` might not be needed at all? Mypy still runs fine on this file without additional error if I remove that part ########## 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: And should this manual stage also be run in the CI build that runs pre-commit? (that is maybe a bit duplicative with the `python_test_type_annotations.sh` script, but it seems important that the pre-commit one works since that is recommended for local usage) ########## python/pyarrow/tests/test_convert_builtin.py: ########## @@ -2452,6 +2457,7 @@ def test_nested_auto_chunking(ty, char): } [email protected] Review Comment: This change should not matter for type checking? But it seems to point to another issue, that we are not actually running those tests without numpy? (if so, I would leave that out here and open an issue for it instead) ########## python/pyproject.toml: ########## Review Comment: According to https://mypy.readthedocs.io/en/stable/stubs.html#creating-a-stub, separate stub packages have to be installed (or on the path?), and don't work with mypy_path ########## python/pyproject.toml: ########## Review Comment: ```suggestion ``` The `files` above already ensures the pyarrow repo is not yet looked at by default when running just `mypy`, but removing it here ensure you can at least already try `mypy pyarrow/tests` locally. ########## python/pyarrow-stubs/pyarrow/__lib/error.pyi: ########## @@ -0,0 +1,71 @@ +# 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 types import TracebackType + +from typing_extensions import Self Review Comment: ```suggestion from typing import Self ``` ########## 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 Review Comment: (and the same in some other files) ########## 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: Additionally, because pandas is being imported with `# type: ignore[import-not-found,import-untyped]`, this `pd.Series` actually becomes `Any`, making the type hint for `mask` essentially useless (it accepts anything). (I was trying to add a small test to `python/pyarrow-stubs/tests/test_array.py` to verify how it catches wrong `mask` arguments, and couldn't let mypy trigger an error) ########## .pre-commit-config.yaml: ########## @@ -136,15 +136,101 @@ repos: - id: flake8 alias: python name: Python Lint + additional_dependencies: + - flake8-pyi args: - "--config" - "python/setup.cfg" + - "--extend-select" + - "Y" + - "--per-file-ignores" + - "python/pyarrow-stubs/pyarrow/*.pyi:E301,E302,E305,E701,F821" files: >- ^(c_glib|dev|python)/ + types: + - file + types_or: + - python + - pyi exclude: >- ( ?^python/pyarrow/vendored/| ) + # Linting, formatting and type checking of the type stubs. These hooks only + # run on request, not on commit or in the existing CI jobs: + # pre-commit run --hook-stage manual python-stubs --all-files + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.17.2 + hooks: + - id: ruff-check + alias: python-stubs + name: Python Stubs Lint (ruff) + stages: + - manual Review Comment: ruff does not have to be manual? (it is fast, installs easily isolated) ########## 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 Review Comment: Can be cleaned up (we already require python 3.11) ########## python/pyarrow-stubs/pyarrow/_stubs_typing.pyi: ########## @@ -0,0 +1,116 @@ +# 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 datetime as dt + +from _typeshed import ( + SupportsDunderGE, + SupportsDunderGT, + SupportsDunderLE, + SupportsDunderLT, +) +from collections.abc import Collection, Container, Iterator, Sequence, Sized +from decimal import Decimal +from typing import Any, Literal, Protocol, TypeAlias + +import numpy as np + +from numpy.typing import NDArray + +from pyarrow import lib +from pyarrow.lib import ChunkedArray + +ArrayLike: TypeAlias = Any +ScalarLike: TypeAlias = Any +Order: TypeAlias = Literal["ascending", "descending"] +JoinType: TypeAlias = Literal[ + "left semi", + "right semi", + "left anti", + "right anti", + "inner", + "left outer", + "right outer", + "full outer", +] +Compression: TypeAlias = Literal[ + "gzip", "bz2", "brotli", "lz4", "lz4_frame", "lz4_raw", "zstd", "snappy" +] +MapsAsPydicts: TypeAlias = Literal["lossy", "strict"] +NullEncoding: TypeAlias = Literal["mask", "encode"] +NullPlacement: TypeAlias = Literal["at_start", "at_end"] +NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] +TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] + +PyScalar: TypeAlias = ( + bool + | int + | float + | Decimal + | str + | bytes + | dt.date + | dt.datetime + | dt.time + | dt.timedelta +) +NumpyScalar: TypeAlias = np.generic[Any] + +Mask: TypeAlias = ( + Sequence[bool | None] | NDArray[np.bool_] | lib.BooleanArray | ChunkedArray Review Comment: ChunkedArray is not actually accepted as a mask in eg pa.array() Also, I think this should be typed as Array instead of BooleanArray, to avoid errors for users if they have not narrowed it down. (or just leave this as `Sequence` for now, and improve this in a follow-up) ########## 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: ... Review Comment: This just always returns `Array`? (same for `repeat` below) ########## 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 Review Comment: Can we be consistent with how things are imported? Above some other names are imported from `pyarrow.lib`. Or is there a reason that it has to be different? (I don't know if circular import issues also play a role in stubs) -- 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]
