GH-49103: [Python] Add internal type system stubs (_types, error, _stubs_typing)

.pre-commit-config.yaml

@@ -136,11 +136,22 @@ 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"
         files: >-
           ^(c_glib|dev|python)/
+        types:
+          - file
+        types_or:
+          - python
+          - pyi
         exclude: >-
           (
           ?^python/pyarrow/vendored/|

docs/source/developers/python/building.rst

@@ -204,6 +204,7 @@ Windows tab under the :ref:`pyarrow_build_section` section.
 
          $ export ARROW_HOME=$(pwd)/dist
          $ export LD_LIBRARY_PATH=$(pwd)/dist/lib:$LD_LIBRARY_PATH
+         $ export DYLD_LIBRARY_PATH=$(pwd)/dist/lib:$DYLD_LIBRARY_PATH
          $ export CMAKE_PREFIX_PATH=$ARROW_HOME:$CMAKE_PREFIX_PATH
 
    .. tab-item:: Windows

python/CMakeLists.txt

@@ -927,6 +927,11 @@ if(CYTHON_VERSION VERSION_GREATER_EQUAL "3.1.0a0")
   list(APPEND CYTHON_FLAGS "-Xfreethreading_compatible=True")
 endif()
 
+set(PYARROW_EDITABLE_BUILD OFF)
+if(SKBUILD_STATE STREQUAL "editable")
+  set(PYARROW_EDITABLE_BUILD ON)
+endif()
+
 foreach(module ${CYTHON_EXTENSIONS})
   string(REPLACE "." ";" directories ${module})
   list(GET directories -1 module_name)
@@ -968,6 +973,11 @@ foreach(module ${CYTHON_EXTENSIONS})
         continue()
       endif()
     endif()
+    if(PYARROW_EDITABLE_BUILD
+       AND module_name STREQUAL "lib"
+       AND output STREQUAL "lib.h")
+      continue()
+    endif()
     install(FILES ${CMAKE_CURRENT_BINARY_DIR}/${output} DESTINATION ".")
   endforeach()
 endforeach()
@@ -1033,10 +1043,17 @@ endif()
 # alongside the package so type checkers can find them (PEP 561).
 set(PYARROW_STUBS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/pyarrow-stubs/pyarrow")
 if(EXISTS "${PYARROW_STUBS_SOURCE_DIR}")
+  set(PYARROW_STUB_INSTALL_ARGS FILES_MATCHING PATTERN "*.pyi")
+  if(PYARROW_EDITABLE_BUILD)
+    list(APPEND
+         PYARROW_STUB_INSTALL_ARGS
+         PATTERN
+         "lib.pyi"
+         EXCLUDE)
+  endif()
   install(DIRECTORY "${PYARROW_STUBS_SOURCE_DIR}/"
           DESTINATION "."
-          FILES_MATCHING
-          PATTERN "*.pyi")
+          ${PYARROW_STUB_INSTALL_ARGS})
 
   if(PYARROW_REQUIRE_STUB_DOCSTRINGS)
     install(CODE "

python/pyarrow-stubs/pyarrow/__init__.pyi

@@ -15,11 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Type stubs for PyArrow.
-
-This is a placeholder stub file.
-Complete type annotations will be added in subsequent PRs.
-"""
+# Type stubs for PyArrow.
+# This is a placeholder stub file.
+# Complete type annotations will be added in subsequent PRs.
 
 from typing import Any
 

python/pyarrow-stubs/pyarrow/_stubs_typing.pyi

@@ -0,0 +1,154 @@
+# 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 collections.abc import Collection, Container, Iterator, Sequence, Sized
+from decimal import Decimal
+from typing import Any, Literal, Protocol, TypeAlias, TypeVar
+
+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"
+]
+NullEncoding: TypeAlias = Literal["mask", "encode"]
+NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"]
+TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"]
+
+IntegerType: TypeAlias = (
+    lib.Int8Type
+    | lib.Int16Type
+    | lib.Int32Type
+    | lib.Int64Type
+    | lib.UInt8Type
+    | lib.UInt16Type
+    | lib.UInt32Type
+    | lib.UInt64Type
+)
+
+Mask: TypeAlias = (
+    Sequence[bool | None]
+    | NDArray[np.bool_]
+    | lib.Array[lib.Scalar[lib.BoolType]]
+    | ChunkedArray[Any]
+)
+Indices: TypeAlias = (
+    Sequence[int | None]
+    | NDArray[np.integer[Any]]
+    | lib.Array[lib.Scalar[IntegerType]]
+    | ChunkedArray[Any]
+)
+
+PyScalar: TypeAlias = (
+    bool
+    | int
+    | float
+    | Decimal
+    | str
+    | bytes
+    | dt.date
+    | dt.datetime
+    | dt.time
+    | dt.timedelta
+)
+
+_T = TypeVar("_T")
+_V = TypeVar("_V", covariant=True)
+
+SingleOrList: TypeAlias = list[_T] | _T
+
+
+class SupportsDunderEQ(Protocol):
+    def __eq__(self, other: object, /) -> bool: ...
+
+
+class SupportsDunderLT(Protocol):
+    def __lt__(self, other: object, /) -> bool: ...
+
+
+class SupportsDunderGT(Protocol):
+    def __gt__(self, other: object, /) -> bool: ...
+
+
+class SupportsDunderLE(Protocol):
+    def __le__(self, other: object, /) -> bool: ...
+
+
+class SupportsDunderGE(Protocol):
+    def __ge__(self, other: object, /) -> bool: ...
+
+
+FilterTuple: TypeAlias = (
+    tuple[str, Literal["=", "==", "!="], SupportsDunderEQ]
+    | tuple[str, Literal["<"], SupportsDunderLT]
+    | tuple[str, Literal[">"], SupportsDunderGT]
+    | tuple[str, Literal["<="], SupportsDunderLE]
+    | tuple[str, Literal[">="], SupportsDunderGE]
+    | tuple[str, Literal["in", "not in"], Collection]
+    | tuple[str, str, Any]  # Allow general str for operator to avoid type errors
+)
+
+
+class Buffer(Protocol): ...
+
+
+class SupportsPyBuffer(Protocol): ...
+
+
+class SupportsArrowStream(Protocol):
+    def __arrow_c_stream__(self, requested_schema=None, /) -> Any: ...
+
+
+class SupportsPyArrowArray(Protocol):
+    def __arrow_array__(self, type=None, /) -> Any: ...
+
+
+class SupportsArrowArray(Protocol):
+    def __arrow_c_array__(self, requested_schema=None, /) -> Any: ...
+
+
+class SupportsArrowDeviceArray(Protocol):
+    def __arrow_c_device_array__(self, requested_schema=None, /, **kwargs) -> Any: ...
+
+
+class SupportsArrowSchema(Protocol):
+    def __arrow_c_schema__(self) -> Any: ...
+
+
+class NullableCollection(Sized, Container[_V], Protocol[_V]):
+    def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ...
+    def __len__(self) -> int: ...
+    def __contains__(self, item: Any, /) -> bool: ...