GH-55263: [Python][Annotations] Add internal types and core data structures

.github/workflows/dev.yml

@@ -67,6 +67,12 @@ jobs:
       - name: Run pre-commit
         run: |
           pre-commit run --all-files --color=always --show-diff-on-failure
+      - name: Run pre-commit for type stubs
+        env:
+          # The previous step already ran the ruff hooks.
+          SKIP: ruff-check,ruff-format
+        run: |
+          pre-commit run --hook-stage manual python-stubs --all-files --color=always --show-diff-on-failure
 
   release:
     name: Source Release and Merge Script on ${{ matrix.runs-on }}

.pre-commit-config.yaml

@@ -145,6 +145,70 @@ repos:
           (
           ?^python/pyarrow/vendored/|
           )
+  # Linting, formatting and type checking of the type stubs. The type checkers
+  # don't run on commit, only on request (the Lint CI job runs them too):
+  #   pre-commit run --hook-stage manual python-stubs --all-files
+  # The type checkers run with uvx, in the environment of
+  # python/requirements-stubs.txt. They don't need PyArrow installed: the
+  # modules that aren't annotated yet, e.g. pyarrow.compute, are treated as Any
+  # (see python/pyproject.toml).
+  - 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)
+        args:
+          - "--fix"
+        files: >-
+          ^python/pyarrow-stubs/
+        types:
+          - pyi
+      - id: ruff-format
+        alias: python-stubs
+        name: Python Stubs Format (ruff)
+        files: >-
+          ^python/pyarrow-stubs/
+        types:
+          - pyi
+  - repo: local
+    hooks:
+      - id: mypy
+        alias: python-stubs
+        name: Python Stubs Type Check (mypy)
+        stages:
+          - manual
+        language: python
+        entry: bash -c "cd python && exec uvx --with-requirements requirements-stubs.txt mypy"
+        pass_filenames: false
+        files: >-
+          ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$)
+        additional_dependencies: &uv
+          - uv==0.13.7
+      - id: pyright
+        alias: python-stubs
+        name: Python Stubs Type Check (pyright)
+        stages:
+          - manual
+        language: python
+        entry: uvx --with-requirements python/requirements-stubs.txt pyright -p python/pyproject.toml
+        pass_filenames: false
+        files: >-
+          ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$)
+        additional_dependencies: *uv
+      - id: ty
+        alias: python-stubs
+        name: Python Stubs Type Check (ty)
+        stages:
+          - manual
+        language: python
+        # Unset the VIRTUAL_ENV of the hook, so that ty uses the environment
+        # of uvx.
+        entry: env -u VIRTUAL_ENV uvx --with-requirements python/requirements-stubs.txt ty check --project python
+        pass_filenames: false
+        files: >-
+          ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$)
+        additional_dependencies: *uv
   - repo: https://github.com/MarcoGorelli/cython-lint
     rev: v0.21.1
     hooks:

ci/scripts/python_test_type_annotations.sh

@@ -25,14 +25,20 @@ if [ -n "${ARROW_PYTHON_VENV:-}" ]; then
   . "${ARROW_PYTHON_VENV}/bin/activate"
 fi
 
-# Install library stubs. Note some libraries contain their own type hints so they need to be installed.
-pip install fsspec pandas-stubs scipy-stubs types-cffi types-psutil types-requests types-python-dateutil
+cd "${pyarrow_dir}"
 
-# Install type checkers
-pip install mypy pyright ty
+# Install type checkers and library stubs. Note some libraries contain their
+# own type hints so they need to be installed.
+pip install -r requirements-stubs.txt
 
 # Run type checkers
-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 expected annotated types against the runtime types.
+pytest -r s pyarrow-stubs-tests
+
+# Check that the stubs match the signatures of the installed PyArrow.
+python scripts/run_stubtest.py

ci/scripts/python_wheel_validate_contents.py

@@ -60,7 +60,7 @@ def validate_wheel(path):
         ), "pyarrow/py.typed is present in the wheel."
 
         source_root = Path(__file__).resolve().parents[2]
-        stubs_dir = source_root / "python" / "pyarrow-stubs" / "pyarrow"
+        stubs_dir = source_root / "python" / "pyarrow-stubs"
         assert stubs_dir.exists(), f"Stub source directory not found: {stubs_dir}"
 
         expected_stub_files = {
@@ -75,7 +75,7 @@ def validate_wheel(path):
         }
 
         assert not (wheel_stub_files == expected_stub_files), (
-            "Wheel .pyi files do not differ from python/pyarrow-stubs/pyarrow.\n"
+            "Wheel .pyi files do not differ from python/pyarrow-stubs.\n"
             f"Missing in wheel: {sorted(expected_stub_files - wheel_stub_files)}\n"
             f"Unexpected in wheel: {sorted(wheel_stub_files - expected_stub_files)}"
         )

docs/source/developers/python/development.rst

@@ -105,15 +105,41 @@ Type Checking
 =============
 
 PyArrow provides type stubs (``*.pyi`` files) for static type checking. These
-stubs are located in the ``pyarrow-stubs/`` directory and are automatically
-included in the distributed wheel packages.
+stubs are located in the ``pyarrow-stubs/`` directory, which is laid out as a
+:pep:`561` stub-only package for ``pyarrow``, so type checkers and IDEs pick
+it up from the ``arrow/python`` directory. The stubs aren't included in the
+distributed wheel packages until they are complete (see
+`GH-49831 <https://github.com/apache/arrow/issues/49831>`_).
 
 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 type checkers and stub packages pinned in ``requirements-stubs.txt`` (with
+uv). This also lints and formats the stubs with ruff:
+
+.. code-block::
+
+   $ pre-commit run --hook-stage manual python-stubs --all-files
+
+The type checkers only run on request, while ruff also runs on commit. The
+Lint CI job runs all of them as well. CI also runs the type checkers, the type
+stub tests and stubtest (see below) against the PyArrow it builds, with
+``ci/scripts/python_test_type_annotations.sh``.
+
+To run a type checker directly instead, install the pinned type checkers and
+stub packages in your PyArrow development environment:
+
+.. code-block::
+
+   $ pip install -r requirements-stubs.txt
+
+Checking the PyArrow tests also needs their dependencies from
+``requirements-test.txt``.
+
 **mypy**
 
 To run mypy on the PyArrow codebase:
@@ -124,6 +150,12 @@ To run mypy on the PyArrow codebase:
    $ mypy
 
 The mypy configuration is in the ``[tool.mypy]`` section of ``pyproject.toml``.
+mypy only uses the stubs when it checks them too, so pass ``pyarrow-stubs``
+along with any other files you check:
+
+.. code-block::
+
+   $ mypy pyarrow-stubs pyarrow/tests/test_array.py
 
 **pyright**
 
@@ -145,6 +177,20 @@ To run ty (note: currently only partially configured):
    $ cd arrow/python
    $ ty check
 
+**stubtest**
+
+mypy's stubtest compares the stubs with the signatures of the installed
+PyArrow. Run it with:
+
+.. code-block::
+
+   $ cd arrow/python
+   $ python scripts/run_stubtest.py
+
+Stubs that aren't written yet are skipped. The known mismatches are listed in
+``pyarrow-stubs-tests/stubtest_allowlist.txt``; remove the entries of the
+mismatches you fix, as stubtest fails on unused entries.
+
 Maintaining Type Stubs
 -----------------------
 
@@ -169,6 +215,44 @@ 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
+
+   from typing import TYPE_CHECKING
+
+   import pyarrow as pa
+
+   from _check import assert_type, check
+
+   def test_list_type() -> None:
+       lst = pa.list_(pa.int32())
+       check(assert_type(lst, pa.ListType), pa.ListType)
+       check(assert_type(lst.value_type, pa.DataType), pa.DataType)
+
+       if TYPE_CHECKING:
+           pa.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
+
+Code under ``if TYPE_CHECKING:`` is never executed. It holds the usage that
+should fail type checking, which is asserted with a
+``# type: ignore[<mypy code>]`` comment: unused ignores are reported as errors
+in these tests, so the test fails once the usage is no longer rejected. ty
+doesn't honor mypy error codes, so it doesn't check invalid usage.
+
 Doctest
 =======
 

python/CMakeLists.txt

@@ -1077,15 +1077,15 @@ endif()
 #
 # TODO(GH-32609): Reintroduce type stubs into PyArrow wheels.
 # GH-49831: Temporarily do not install pyarrow-stubs into wheels.
-# Stubs live in pyarrow-stubs/pyarrow/ during development but are installed
+# Stubs live in pyarrow-stubs/ during development but are installed
 # alongside the package so type checkers can find them (PEP 561).
 # The stubs are currently incomplete, and some type checkers consume .pyi files
 # even without the py.typed marker. Re-enable this when the stubs are complete.
 if(PYARROW_REQUIRE_STUB_DOCSTRINGS)
   message(FATAL_ERROR "PYARROW_REQUIRE_STUB_DOCSTRINGS cannot be used while "
                       "pyarrow-stubs are omitted from wheels (GH-49831).")
 endif()
-# set(PYARROW_STUBS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/pyarrow-stubs/pyarrow")
+# set(PYARROW_STUBS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/pyarrow-stubs")
 # if(EXISTS "${PYARROW_STUBS_SOURCE_DIR}")
 #   install(DIRECTORY "${PYARROW_STUBS_SOURCE_DIR}/"
 #           DESTINATION "."

python/pyarrow-stubs-tests/_check.py

@@ -0,0 +1,48 @@
+# 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.list_(pa.int32()), pa.ListType), pa.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:``
+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.
+"""
+
+from typing import Any, TypeVar, assert_type
+
+__all__ = ["assert_type", "check"]
+
+_T = TypeVar("_T")
+
+
+def check(actual: _T, klass: type[Any] | tuple[type[Any], ...]) -> _T:
+    """Assert that ``actual`` is an instance of ``klass`` and return it."""
+    if not isinstance(actual, klass):
+        raise AssertionError(
+            f"Expected an instance of {klass}, got {type(actual)}: {actual!r}"
+        )
+    return actual

python/pyarrow-stubs-tests/conftest.py

@@ -0,0 +1,37 @@
+# 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
+import types
+
+import pyarrow as pa
+
+
+def _stub_only_type(name: str) -> type[pa.DataType]:
+    # The factories of the stub-only DataType subclasses (Int32Type, ...)
+    # return plain DataType instances at runtime.
+    if name.endswith("Type") and not name.startswith("_"):
+        return pa.DataType
+    raise AttributeError(name)
+
+
+# pyarrow._stubs_typing only exists in the type stubs. Provide it at runtime,
+# so that the tests can import the stub-only types without a TYPE_CHECKING
+# block.
+_stubs_typing = types.ModuleType("pyarrow._stubs_typing")
+setattr(_stubs_typing, "__getattr__", _stub_only_type)
+sys.modules["pyarrow._stubs_typing"] = _stubs_typing

python/pyarrow-stubs-tests/stubtest_allowlist.txt

@@ -0,0 +1,30 @@
+# 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.
+
+# Mismatches between the type stubs and the runtime that stubtest skips, one
+# regex per line; see scripts/run_stubtest.py. Unused entries fail stubtest,
+# so remove the entries of the mismatches you fix. Print the entries for the
+# current mismatches with
+#   python scripts/run_stubtest.py --generate-allowlist
+
+# Stub-only modules. The pyarrow.__lib modules are checked against
+# pyarrow.lib, which has no __all__.
+pyarrow\._stubs_typing
+pyarrow\.__lib\.[a-z]+\.__all__
+
+# The __init__ of Cython extension types takes (*args, **kwargs) at runtime.
+pyarrow\.__lib\.[a-z]+\.[A-Za-z_]+\.__init__