review feedback

Author: rok

ci/scripts/python_test_type_annotations.sh

@@ -20,24 +20,20 @@
 set -ex
 pyarrow_dir=${1}
 
-if [ "${PYARROW_TEST_ANNOTATIONS}" == "ON" ]; then
-  if [ -n "${ARROW_PYTHON_VENV:-}" ]; then
-    # shellcheck source=/dev/null
-    . "${ARROW_PYTHON_VENV}/bin/activate"
-  fi
+if [ -n "${ARROW_PYTHON_VENV:-}" ]; then
+  # shellcheck source=/dev/null
+  . "${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 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
+# 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
+# Run type checkers
+pushd "${pyarrow_dir}"
+mypy
+pyright
+ty check
+popd

python/pyproject.toml

@@ -113,7 +113,6 @@ exclude = [
 # TODO: Enable type checking once stubs are merged
 [tool.pyright]
 pythonPlatform = "All"
-pythonVersion = "3.10"
 include = ["pyarrow-stubs"]
 exclude = [
     "pyarrow",

python/requirements-wheel-build.txt

@@ -1,5 +1,7 @@
 build
 cython>=3.1
+# Needed for build-time stub docstring extraction
+libcst>=1.8.6
 numpy>=2.0.0
 setuptools_scm
 setuptools>=77

dev/update_stub_docstrings.py python/scripts/update_stub_docstrings.pydev/update_stub_docstrings.py renamed to python/scripts/update_stub_docstrings.py

@@ -19,13 +19,12 @@
 Extract docstrings from pyarrow runtime and insert them into stub files.
 
 Usage (from python/ directory with pyarrow built):
-    python ../dev/update_stub_docstrings.py pyarrow-stubs
+    python scripts/update_stub_docstrings.py pyarrow-stubs
 """
 
 import argparse
 import importlib
 import inspect
-import shutil
 import sys
 from pathlib import Path
 from textwrap import indent
@@ -199,27 +198,19 @@ def add_docstrings_to_stubs(stubs_dir):
         stub_file.write_text(modified.code)
 
 
-def copy_stubs(src_dir, dest_dir):
-    """Copy .pyi files from src_dir to dest_dir."""
-    src_dir, dest_dir = Path(src_dir), Path(dest_dir)
-    if not src_dir.exists():
-        return
+def add_docstrings_from_build(stubs_dir, build_lib):
+    """
+    Entry point for setup.py: update docstrings using pyarrow from build directory.
 
-    print(f"Copying stubs: {src_dir} -> {dest_dir}")
-    for src in src_dir.rglob('*.pyi'):
-        dest = dest_dir / src.relative_to(src_dir)
-        dest.parent.mkdir(parents=True, exist_ok=True)
-        shutil.copy2(src, dest)
-
-
-def update_stubs_for_build(stubs_dir, build_lib):
-    """Entry point for setup.py: update docstrings and copy stubs to build dir."""
+    During the build process, pyarrow is not installed in the system Python.
+    We need to temporarily add the build directory to sys.path so we can
+    import pyarrow and extract docstrings from it.
+    """
     stubs_dir, build_lib = Path(stubs_dir), Path(build_lib)
 
     sys.path.insert(0, str(build_lib))
     try:
         add_docstrings_to_stubs(stubs_dir)
-        copy_stubs(stubs_dir / "pyarrow", build_lib / "pyarrow")
     finally:
         sys.path.pop(0)
 
@@ -229,5 +220,9 @@ def update_stubs_for_build(stubs_dir, build_lib):
     parser.add_argument("stubs_dir", type=Path, help="Path to pyarrow-stubs folder")
     args = parser.parse_args()
 
-    sys.path.insert(0, ".")
+    # Add the directory containing this script's parent (python/) to sys.path
+    # so pyarrow can be imported when running from the python/ directory
+    script_dir = Path(__file__).resolve().parent
+    python_dir = script_dir.parent
+    sys.path.insert(0, str(python_dir))
     add_docstrings_to_stubs(args.stubs_dir.resolve())

python/setup.py

@@ -133,16 +133,35 @@ def _update_stubs(self):
         build_cmd = self.get_finalized_command('build')
         build_lib = os.path.abspath(build_cmd.build_lib)
 
-        # Import here to avoid hard dependency on the dev script
-        sys.path.insert(0, pjoin(setup_dir, '..', 'dev'))
+        # Import the stub docstring updater from python/scripts
+        sys.path.insert(0, pjoin(setup_dir, 'scripts'))
         try:
-            from update_stub_docstrings import update_stubs_for_build
-            update_stubs_for_build(stubs_dir, build_lib)
-        except ImportError:
-            print("-- Skipping stubs (update_stub_docstrings.py not found)")
+            from update_stub_docstrings import add_docstrings_from_build
+            add_docstrings_from_build(stubs_dir, build_lib)
         finally:
             sys.path.pop(0)
 
+        # Copy stub files to build directory
+        self._copy_stubs(stubs_dir, build_lib)
+
+    def _copy_stubs(self, stubs_dir, build_lib):
+        """Copy .pyi stub files to the build directory."""
+        src_dir = pjoin(stubs_dir, 'pyarrow')
+        dest_dir = pjoin(build_lib, 'pyarrow')
+
+        if not os.path.exists(src_dir):
+            return
+
+        print(f"-- Copying stubs: {src_dir} -> {dest_dir}")
+        for root, dirs, files in os.walk(src_dir):
+            for fname in files:
+                if fname.endswith('.pyi'):
+                    src = pjoin(root, fname)
+                    rel_path = os.path.relpath(src, src_dir)
+                    dest = pjoin(dest_dir, rel_path)
+                    os.makedirs(os.path.dirname(dest), exist_ok=True)
+                    shutil.copy2(src, dest)
+
     # adapted from cmake_build_ext in dynd-python
     # github.com/libdynd/dynd-python