Generate pylibcudf stubs with stubgen-pyx by vyasr · Pull Request #23628 · NVIDIA/cudf

vyasr · 2026-08-12T17:10:17Z

Summary

add a local pre-commit hook that regenerates tracked public pylibcudf .pyi files with stubgen-pyx==0.2.18
add a wrapper script to scope generation to the existing public pylibcudf stub set and preserve SPDX headers
replace the handwritten pylibcudf stubs with generated output for review

Motivation

We have been improving stubgen-pyx upstream, and this draft gives cudf developers a concrete diff for evaluating whether generated pylibcudf stubs are close enough to replace the handwritten files.

Notes for reviewers

This is intentionally a draft/evaluation PR. The generated stubs are much more complete and source-derived than the handwritten versions, but there are still visible gaps worth reviewing:

CudaStreamLike is lost in generated signatures, generally becoming object | None
some curated public helper aliases and protocols disappear in favor of lower-level generated imports
some generated signatures expose implementation details or private helpers that the handwritten stubs hid
full mypy currently reports generated-stub issues and downstream precision regressions

Validation

deleted tracked python/pylibcudf/pylibcudf/**/*.pyi files and regenerated them with pre-commit run stubgen-pyx-pylibcudf --all-files
reran pre-commit run stubgen-pyx-pylibcudf --all-files after rebasing on latest upstream/main; it passed with no further changes
pre-commit run ruff --files ci/checks/generate_pylibcudf_stubs.py .pre-commit-config.yaml
pre-commit run ruff-format --files ci/checks/generate_pylibcudf_stubs.py .pre-commit-config.yaml
pre-commit run yamllint --files .pre-commit-config.yaml
pre-commit run verify-copyright --files ...
git diff --check

Full commit-time mypy was also run while preparing this branch and currently fails on the generated stubs; that failure is part of the review signal for this draft.

copy-pr-bot · 2026-08-12T17:10:20Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

wence- · 2026-08-14T08:59:15Z

Can you temporarily elide the insertion of docstrings into the stubs. That will make the diffs of the types, which is the important thing, much easier to read.

wence- · 2026-08-14T08:57:52Z

+from rmm.pylibrmm.memory_resource import DeviceMemoryResource

-__all__ = ["AvroReaderOptions", "AvroReaderOptionsBuilder", "read_avro"]
+__all__ = ['read_avro', 'AvroReaderOptions', 'AvroReaderOptionsBuilder']


The script should run ruff check/fmt on the pyis.

Good idea, I'll add that.

wence- · 2026-08-14T09:00:27Z

+    def build(self) -> AvroReaderOptions:
+        """Create a AvroReaderOptions object"""
+
+def read_avro(options: AvroReaderOptions, stream: CudaStreamLike | None=None, mr: DeviceMemoryResource | None=None) -> TableWithMetadata:


nit: Usually type stubs don't advertise the default, but rather just show that there is a default by saying mr : SomeType | None = ....

Good point, I'll see if we can modify the generator to do that. Opened jon-edward/stubgen-pyx#62

wence- · 2026-08-14T09:01:22Z

-    TableWithMetadata,
-)
+from pylibcudf.io.types import SinkInfo, SourceInfo, TableWithMetadata
+from pylibcudf.libcudf.io.types import compression_type, quote_style


Need to do renaming of these into CamelCase

Yeah, I haven't yet figured out how best to do that. I made this work locally by supporting an import remapping option for specific type names, which is a bit manual but might be our best option. I don't know how we can get this to work automagically without some input.

wence- · 2026-08-14T09:08:30Z

-        stream: CudaStreamLike | None = None,
-        mr: DeviceMemoryResource | None = None,
-    ) -> Scalar: ...
+    def from_arrow(obj: ArrowLike, dtype: DataType | None=None, stream: CudaStreamLike | None=None, mr: DeviceMemoryResource | None=None) -> ArrowLike:


Return type is wrong.

wence- · 2026-08-14T09:10:09Z

+def _infer_list_depth_and_dtype(obj: list) -> tuple[int, type]:
+    """Infer the nesting depth and final scalar type."""
+def _flatten_nested_list(obj: list, depth: int) -> tuple[list, tuple[int, ...]]:
+    """Flatten a nested list and compute the shape"""
+def _infer_shape(obj: list, depth: int) -> tuple[int, ...]: ...
+def _flatten(obj: list, out: list, offset: int) -> int: ...
+def _prepare_array_metadata(iface: dict) -> tuple[int, int, tuple[int, ...], tuple[int, ...] | None, DataType]:


Let's just cull all underscore names from the type stubs.

jon-edward/stubgen-pyx#63

wence- · 2026-08-14T12:15:25Z

+from pylibcudf.libcudf.types import interpolation as Interpolation
+from pylibcudf.libcudf.types import mask_state as MaskState
+from pylibcudf.libcudf.types import nan_equality as NanEquality
+from pylibcudf.libcudf.types import nan_policy as NanPolicy
+from pylibcudf.libcudf.types import null_aware as NullAware
+from pylibcudf.libcudf.types import null_equality as NullEquality
+from pylibcudf.libcudf.types import null_order as NullOrder
+from pylibcudf.libcudf.types import null_policy as NullPolicy
+from pylibcudf.libcudf.types import order as Order
+from pylibcudf.libcudf.types import output_nullability as OutputNullability
+from pylibcudf.libcudf.types import size_type
+from pylibcudf.libcudf.types import sorted as Sorted
+from pylibcudf.libcudf.types import type_id
+from pylibcudf.libcudf.types import type_id as TypeId
+from pylibcudf.libcudf.utilities.type_dispatcher import type_to_id


None of this will work since these modules don't exist.

Same fundamental issue as #23628 (comment).

vyasr · 2026-08-18T00:26:04Z

Can you temporarily elide the insertion of docstrings into the stubs. That will make the diffs of the types, which is the important thing, much easier to read.

Probably not just temporarily, we should change stubgen-pyx to support eliding docstrings since there's no real reason to keep docstrings in the stubs if they're in the source. Opened jon-edward/stubgen-pyx#64

github-actions Bot assigned vyasr Aug 12, 2026

github-actions Bot added Python Affects Python cuDF API. pylibcudf Issues specific to the pylibcudf package labels Aug 12, 2026

github-project-automation Bot added this to cuDF Python and cuDF Python Aug 12, 2026

github-project-automation Bot moved this to Todo in cuDF Python Aug 12, 2026

vyasr force-pushed the codex/stubgen-pyx-pylibcudf-stubs branch from f7f07d5 to 29f2241 Compare August 13, 2026 20:24

Generate pylibcudf stubs with stubgen-pyx

e748f52

vyasr force-pushed the codex/stubgen-pyx-pylibcudf-stubs branch from 29f2241 to e748f52 Compare August 13, 2026 21:37

vyasr added 2 commits August 14, 2026 14:59

Add pylibcudf return annotations for stub generation

ebd5ea6

Document pylibcudf stub generation guidelines

259d64d

wence- reviewed Aug 17, 2026

View reviewed changes

Uh oh!

Conversation

vyasr commented Aug 12, 2026

Summary

Motivation

Notes for reviewers

Validation

Uh oh!

copy-pr-bot Bot commented Aug 12, 2026

Uh oh!

wence- commented Aug 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vyasr commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

vyasr commented Aug 18, 2026 •

edited

Loading