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

jorisvandenbossche · 2026-11-11T09:20:17Z

That is not what I see locally. Running eg mypy pyarrow/tests/test_table.py works just fine and finds the same issues whether I add pyarrow-stubs or not.
(assuming I pass --follow-untyped-imports --check-untyped-defs, otherwise it finds nothing related to pyarrow anyhow)

-Original file line number
+Diff line change
@@ Expand Up / @@ -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 }}
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -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:
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -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

-Original file line number
+Diff line change
@@ Expand Up / @@ -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 = {
@@ Expand All / @@ -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)}"
             )
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -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:
@@ Expand All / @@ -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**
@@ Expand All / @@ -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
     -----------------------
@@ Expand All / @@ -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
     =======
@@ Expand Down @@

Navigation Menu

Repository navigation

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

Uh oh!

Diff view

Diff view

There are no files selected for viewing

jorisvandenbossche Nov 11, 2026

Uh oh!

Uh oh!

Uh oh!

-Original file line number
+Diff line change
@@ Expand Up / @@ -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 "."
@@ Expand Down @@

-Original file line number
+Diff line change
@@ -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

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

Are you sure you want to change the base?

Uh oh!

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

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

jorisvandenbossche Nov 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!