GH-49273: [Python] Move stub docstring script to _build_utils with gr…

python/MANIFEST.in

@@ -5,7 +5,6 @@ include ../NOTICE.txt
 global-include CMakeLists.txt
 graft pyarrow
 graft pyarrow-stubs
-include scripts/update_stub_docstrings.py
 graft cmake_modules
 
 global-exclude *.so

python/pyarrow/_build_utils/__init__.py

@@ -0,0 +1,16 @@
+# 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.

python/pyarrow/_build_utils/update_stub_docstrings.py

@@ -0,0 +1,230 @@
+# 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.
+
+"""
+Extract docstrings from pyarrow runtime and insert them into stub files.
+
+Usage (from python/ directory with pyarrow built):
+    python scripts/update_stub_docstrings.py pyarrow-stubs
+"""
+
+import argparse
+import importlib
+import inspect
+import sys
+from pathlib import Path
+from textwrap import indent
+
+import libcst
+from libcst import matchers as m
+
+
+def _resolve_object(module, path):
+    """Resolve an object by dotted path from a module."""
+    if not path:
+        return module, None, module.__name__
+
+    parts = path.split(".")
+    parent = None
+    obj = module
+
+    for part in parts:
+        parent = obj
+        try:
+            obj = getattr(obj, part)
+        except AttributeError:
+            try:
+                obj = vars(parent).get(part)
+                if obj is not None:
+                    continue
+            except TypeError:
+                pass
+            return None, None, None
+
+    return obj, parent, getattr(obj, "__name__", parts[-1])
+
+
+def _get_docstring(name, module, indentation):
+    """Extract and format a docstring for insertion into a stub file."""
+    obj, parent, obj_name = _resolve_object(module, name)
+    if obj is None:
+        print(f"{name} not found in {module.__name__}")
+        return None
+
+    docstring = inspect.getdoc(obj)
+    if not docstring:
+        return None
+
+    # Remove signature prefix
+    parent_name = getattr(parent, "__name__", None) if parent else None
+    if docstring.startswith(obj_name) or (
+        parent_name and docstring.startswith(f"{parent_name}.{obj_name}")
+    ):
+        docstring = "\n".join(docstring.splitlines()[2:])
+
+    # Skip empty docstrings
+    if not docstring.strip():
+        return None
+
+    prefix = "    " * indentation
+    return '"""\n' + indent(docstring + '\n"""', prefix)
+
+
+class DocstringInserter(libcst.CSTTransformer):
+    """CST transformer that inserts docstrings into stub file nodes."""
+
+    def __init__(self, module, namespace):
+        self.module = module
+        self.base_namespace = namespace
+        self.stack = []
+        self.indentation = 0
+
+    def _full_name(self):
+        name = ".".join(self.stack)
+        return f"{self.base_namespace}.{name}" if self.base_namespace else name
+
+    def leave_Module(self, original_node, updated_node):
+        new_body = []
+        clone_matcher = m.SimpleStatementLine(
+            body=[m.Assign(value=m.Call(func=m.Name(value="_clone_signature"))),
+                  m.ZeroOrMore()]
+        )
+        for stmt in updated_node.body:
+            new_body.append(stmt)
+            if m.matches(stmt, clone_matcher):
+                name = stmt.body[0].targets[0].target.value
+                if self.base_namespace:
+                    name = f"{self.base_namespace}.{name}"
+                docstring = _get_docstring(name, self.module, 0)
+                if docstring:
+                    new_body.append(libcst.SimpleStatementLine(
+                        body=[libcst.Expr(value=libcst.SimpleString(docstring))]))
+        return updated_node.with_changes(body=new_body)
+
+    def visit_ClassDef(self, node):
+        self.stack.append(node.name.value)
+        self.indentation += 1
+
+    def leave_ClassDef(self, original_node, updated_node):
+        name = self._full_name()
+        docstring = _get_docstring(name, self.module, self.indentation)
+
+        if docstring:
+            ellipsis_class = m.ClassDef(body=m.IndentedBlock(body=[
+                m.SimpleStatementLine(body=[
+                    m.Expr(m.Ellipsis()), m.ZeroOrMore()]), m.ZeroOrMore()]))
+            func_class = m.ClassDef(body=m.IndentedBlock(
+                body=[m.FunctionDef(), m.ZeroOrMore()]))
+
+            if m.matches(updated_node, ellipsis_class):
+                updated_node = updated_node.deep_replace(
+                    updated_node.body.body[0].body[0].value,
+                    libcst.SimpleString(value=docstring))
+            elif m.matches(updated_node, func_class):
+                docstring_stmt = libcst.SimpleStatementLine(
+                    body=[libcst.Expr(value=libcst.SimpleString(value=docstring))])
+                updated_node = updated_node.with_changes(
+                    body=updated_node.body.with_changes(
+                        body=[docstring_stmt] + list(updated_node.body.body)))
+
+        self.stack.pop()
+        self.indentation -= 1
+        return updated_node
+
+    def visit_FunctionDef(self, node):
+        self.stack.append(node.name.value)
+        self.indentation += 1
+
+    def leave_FunctionDef(self, original_node, updated_node):
+        name = self._full_name()
+        ellipsis_func = m.FunctionDef(
+            body=m.SimpleStatementSuite(body=[m.Expr(m.Ellipsis())]))
+
+        if m.matches(original_node, ellipsis_func):
+            docstring = _get_docstring(name, self.module, self.indentation)
+            if docstring:
+                docstring_stmt = libcst.SimpleStatementLine(
+                    body=[libcst.Expr(value=libcst.SimpleString(value=docstring))])
+                updated_node = updated_node.with_changes(
+                    body=libcst.IndentedBlock(body=[docstring_stmt]))
+
+        self.stack.pop()
+        self.indentation -= 1
+        return updated_node
+
+
+LIB_MODULES = {"array", "builder", "compat", "config", "device", "error", "io",
+               "_ipc", "memory", "pandas_shim", "scalar", "table", "tensor", "_types"}
+
+
+def add_docstrings_to_stubs(stubs_dir):
+    """Update all stub files in stubs_dir with docstrings from pyarrow runtime."""
+    stubs_dir = Path(stubs_dir)
+    print(f"Updating stub docstrings in: {stubs_dir}")
+
+    pyarrow = importlib.import_module("pyarrow")
+
+    for stub_file in stubs_dir.rglob('*.pyi'):
+        if stub_file.name == "_stubs_typing.pyi":
+            continue
+
+        module_name = stub_file.stem
+        if module_name in LIB_MODULES:
+            namespace = "lib"
+        elif stub_file.parent.name in ("parquet", "interchange"):
+            namespace = f"{stub_file.parent.name}.{module_name}"
+        elif module_name == "__init__":
+            namespace = ""
+        else:
+            namespace = module_name
+
+        print(f"  {stub_file.name} -> {namespace or '(root)'}")
+        tree = libcst.parse_module(stub_file.read_text())
+        modified = tree.visit(DocstringInserter(pyarrow, namespace))
+        stub_file.write_text(modified.code)
+
+
+def add_docstrings_from_build(stubs_dir, build_lib):
+    """
+    Entry point for setup.py: update docstrings using pyarrow from build directory.
+
+    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)
+    finally:
+        sys.path.pop(0)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("stubs_dir", type=Path,
+                        help="Path to pyarrow-stubs folder")
+    args = parser.parse_args()
+
+    python_dir = Path(__file__).resolve().parent.parent.parent
+    sys.path.insert(0, str(python_dir))
+    add_docstrings_to_stubs(args.stubs_dir.resolve())
+
+
+if __name__ == "__main__":
+    main()

python/pyproject.toml

@@ -18,8 +18,8 @@
 [build-system]
 requires = [
     "cython >= 3.1",
-    # Needed for build-time stub docstring extraction
-    "libcst>=1.8.6",
+    # Optional: enables stub docstring injection during build
+    # "libcst>=1.8.6",
     "numpy>=1.25",
     # configuring setuptools_scm in pyproject.toml requires
     # versions released after 2022
@@ -87,6 +87,7 @@ include-package-data=true
 
 [tool.setuptools.packages.find]
 include = ["pyarrow"]
+exclude = ["pyarrow._build_utils", "pyarrow._build_utils.*"]
 namespaces = false
 
 [tool.setuptools.package-data]