GH-32609: [CI] Add type checking infrastructure and CI workflow for type annotations

.github/workflows/python.yml

@@ -239,6 +239,11 @@ jobs:
       - name: Test
         shell: bash
         run: ci/scripts/python_test.sh $(pwd) $(pwd)/build
+      - name: Test annotations
+        shell: bash
+        env:
+          PYARROW_TEST_ANNOTATIONS: "ON"
+        run: ci/scripts/python_test_type_annotations.sh $(pwd)/python
 
   windows:
     name: AMD64 Windows 2022 Python 3.13
@@ -296,3 +301,9 @@ jobs:
         shell: cmd
         run: |
           call "ci\scripts\python_test.bat" %cd%
+      - name: Test annotations
+        shell: cmd
+        env:
+          PYARROW_TEST_ANNOTATIONS: "ON"
+        run: |
+          call "ci\scripts\python_test_type_annotations.bat" %cd%\python

ci/scripts/python_test_type_annotations.bat

@@ -0,0 +1,40 @@
+@rem Licensed to the Apache Software Foundation (ASF) under one
+@rem or more contributor license agreements.  See the NOTICE file
+@rem distributed with this work for additional information
+@rem regarding copyright ownership.  The ASF licenses this file
+@rem to you under the Apache License, Version 2.0 (the
+@rem "License"); you may not use this file except in compliance
+@rem with the License.  You may obtain a copy of the License at
+@rem
+@rem   http://www.apache.org/licenses/LICENSE-2.0
+@rem
+@rem Unless required by applicable law or agreed to in writing,
+@rem software distributed under the License is distributed on an
+@rem "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+@rem KIND, either express or implied.  See the License for the
+@rem specific language governing permissions and limitations
+@rem under the License.
+
+@echo on
+
+set PYARROW_DIR=%1
+
+if "%PYARROW_TEST_ANNOTATIONS%"=="ON" (
+    echo Annotation testing on Windows ...
+
+    @REM Install library stubs. Note some libraries contain their own type hints so they need to be installed.
+    %PYTHON_CMD% -m pip install fsspec pandas-stubs scipy-stubs types-cffi types-psutil types-requests types-python-dateutil || exit /B 1
+
+    @REM Install type checkers
+    %PYTHON_CMD% -m pip install mypy pyright ty || exit /B 1
+
+    @REM Run type checkers
+    pushd %PYARROW_DIR%
+
+    mypy || exit /B 1
+    pyright || exit /B 1
+    ty check || exit /B 1
+    popd
+) else (
+    echo Annotation testing skipped on Windows ...
+)

ci/scripts/python_test_type_annotations.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+#
+# 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.
+
+set -ex
+pyarrow_dir=${1}
+
+if [ "${PYARROW_TEST_ANNOTATIONS}" == "ON" ]; then
+  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
+
+  # Install type checkers
+  pip install mypy pyright ty
+
+  # Run type checkers
+  pushd ${pyarrow_dir}
+  mypy
+  pyright
+  ty check;
+  popd
+else
+  echo "Skipping type annotation tests";
+fi

ci/scripts/python_wheel_macos_build.sh

@@ -177,6 +177,11 @@ export CMAKE_PREFIX_PATH=${build_dir}/install
 export SETUPTOOLS_SCM_PRETEND_VERSION=${PYARROW_VERSION}
 
 pushd ${source_dir}/python
+# We first populate stub docstrings and then build the wheel
+python setup.py build_ext --inplace
+python -m pip install libcst
+python ../dev/update_stub_docstrings.py pyarrow-stubs
+
 python setup.py bdist_wheel
 popd
 

ci/scripts/python_wheel_validate_contents.py

@@ -34,7 +34,7 @@ def validate_wheel(path):
     ]
     assert not outliers, f"Unexpected contents in wheel: {sorted(outliers)}"
     print(f"The wheel: {wheels[0]} seems valid.")
-
+    # TODO(GH-32609): Validate some docstrings were generated and added.
 
 def main():
     parser = argparse.ArgumentParser()

ci/scripts/python_wheel_windows_build.bat

@@ -132,6 +132,11 @@ set CMAKE_PREFIX_PATH=C:\arrow-dist
 
 pushd C:\arrow\python
 
+@REM We first populate stub docstrings and then build the wheel
+%PYTHON_CMD% setup.py build_ext --inplace
+%PYTHON_CMD% -m pip install libcst
+%PYTHON_CMD% ..\dev\update_stub_docstrings.py pyarrow-stubs
+
 @REM Build wheel
 %PYTHON_CMD% setup.py bdist_wheel || exit /B 1
 

ci/scripts/python_wheel_xlinux_build.sh

@@ -167,6 +167,11 @@ export ARROW_HOME=/tmp/arrow-dist
 export CMAKE_PREFIX_PATH=/tmp/arrow-dist
 
 pushd /arrow/python
+# We first populate stub docstrings and then build the wheel
+python setup.py build_ext --inplace
+python -m pip install libcst
+python ../dev/update_stub_docstrings.py pyarrow-stubs
+
 python setup.py bdist_wheel
 
 echo "=== Strip symbols from wheel ==="

compose.yaml

@@ -919,12 +919,14 @@ services:
     environment:
       <<: [*common, *ccache, *sccache]
       PYTEST_ARGS:  # inherit
+      PYARROW_TEST_ANNOTATIONS: "ON"
     volumes: *conda-volumes
     command: &python-conda-command
       ["
         /arrow/ci/scripts/cpp_build.sh /arrow /build &&
         /arrow/ci/scripts/python_build.sh /arrow /build &&
-        /arrow/ci/scripts/python_test.sh /arrow"]
+        /arrow/ci/scripts/python_test.sh /arrow &&
+        /arrow/ci/scripts/python_test_type_annotations.sh /arrow/python"]
 
   conda-python-emscripten:
     # Usage:
@@ -1001,14 +1003,16 @@ services:
       ARROW_S3: "OFF"
       ARROW_SUBSTRAIT: "OFF"
       ARROW_WITH_OPENTELEMETRY: "OFF"
+      PYARROW_TEST_ANNOTATIONS: "ON"
       SETUPTOOLS_SCM_PRETEND_VERSION:
     volumes: *ubuntu-volumes
     deploy: *cuda-deploy
     command: &python-command >
       /bin/bash -c "
         /arrow/ci/scripts/cpp_build.sh /arrow /build &&
         /arrow/ci/scripts/python_build.sh /arrow /build &&
-        /arrow/ci/scripts/python_test.sh /arrow"
+        /arrow/ci/scripts/python_test.sh /arrow &&
+        /arrow/ci/scripts/python_test_type_annotations.sh /arrow/python"
 
   debian-python:
     # Usage:
@@ -1500,6 +1504,7 @@ services:
         python: ${PYTHON}
     shm_size: *shm-size
     environment:
+      PYARROW_TEST_ANNOTATIONS: "ON"
       <<: [*common, *ccache, *sccache]
       PARQUET_REQUIRE_ENCRYPTION:  # inherit
       HYPOTHESIS_PROFILE:  # inherit
@@ -1510,7 +1515,8 @@ services:
         /arrow/ci/scripts/cpp_build.sh /arrow /build &&
         /arrow/ci/scripts/python_build.sh /arrow /build &&
         mamba uninstall -y numpy &&
-        /arrow/ci/scripts/python_test.sh /arrow"]
+        /arrow/ci/scripts/python_test.sh /arrow &&
+        /arrow/ci/scripts/python_test_type_annotations.sh /arrow/python"]
 
   conda-python-docs:
     # Usage:
@@ -1530,13 +1536,15 @@ services:
       BUILD_DOCS_CPP: "ON"
       BUILD_DOCS_PYTHON: "ON"
       PYTEST_ARGS: "--doctest-modules --doctest-cython"
+      PYARROW_TEST_ANNOTATIONS: "ON"
     volumes: *conda-volumes
     command:
       ["/arrow/ci/scripts/cpp_build.sh /arrow /build &&
         /arrow/ci/scripts/python_build.sh /arrow /build &&
         pip install -e /arrow/dev/archery[numpydoc] &&
         archery numpydoc --allow-rule GL10,PR01,PR03,PR04,PR05,PR10,RT03,YD01 &&
-        /arrow/ci/scripts/python_test.sh /arrow"]
+        /arrow/ci/scripts/python_test.sh /arrow &&
+        /arrow/ci/scripts/python_test_type_annotations.sh /arrow/python"]
 
   conda-python-dask:
     # Possible $DASK parameters: