Generate pyd type info#2247
Open
themason2011 wants to merge 4 commits into
Open
Conversation
Add pybind11-stubgen based type stub generation for the onnxruntime_genai Python package, including ElementType enum binding, fix_pyi.py post-processing, py.typed marker, and CMake POST_BUILD wiring.
Author
|
@microsoft-github-policy-service agree company="Microsoft" |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds Python typing support to the onnxruntime_genai wheel by generating .pyi stubs from the pybind11 extension at build time (plus a small stub post-processor), and exposes an ElementType enum needed by the public Python surface.
Changes:
- Bind
OgaElementTypeinto Python asElementTypein the pybind11 module. - Add build-time stub generation + stub-fix step (plus
py.typedand an__init__.pyire-export) so the wheel ships type information. - Add
fix_pyi.pyto rewrite generated stubs to use Python-facing names instead of C++/C API names.
Reviewed changes
Copilot reviewed 4 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/python/python.cpp | Exposes ElementType enum via pybind11 so it appears in the Python API (and stubs). |
| src/python/py/py.typed | Adds PEP 561 marker so type checkers treat the package as typed. |
| src/python/fix_pyi.py | Post-processes generated stubs to strip the Oga prefix from type names. |
| src/python/CMakeLists.txt | Adds POST_BUILD stub generation and wires generated/fixed stubs into the wheel layout. |
| src/python/init.pyi.in | Provides a re-exporting package-level stub that mirrors runtime imports. |
Comment on lines
+7
to
+17
| from argparse import ArgumentParser | ||
| from pathlib import Path | ||
|
|
||
| parser = ArgumentParser() | ||
| parser.add_argument("--input", "-i", required=True, help="Input .pyi file") | ||
| args = parser.parse_args() | ||
|
|
||
| pyi_file = Path(args.input) | ||
| content = pyi_file.read_text() | ||
| fixed_content = content.replace("Oga", "") | ||
| pyi_file.write_text(fixed_content) |
Comment on lines
+65
to
+67
| # Be compatible with both system-wide but unprivileged env and virtual env where --user is not available | ||
| COMMAND ${PYTHON_EXECUTABLE} -m pip install --user pybind11-stubgen numpy || ${PYTHON_EXECUTABLE} -m pip install pybind11-stubgen numpy | ||
| COMMAND ${PYTHON_EXECUTABLE} -m pybind11_stubgen $<TARGET_FILE_BASE_NAME:python> --output-dir ${WHEEL_TARGET_NAME} |
Comment on lines
+63
to
+69
| COMMAND ${CMAKE_COMMAND} -E copy_directory ${WHEEL_TARGET_NAME} ${GENERATE_TYPE_DIR} | ||
| COMMAND ${CMAKE_COMMAND} -E copy_directory ${ORT_LIB_DIR} ${GENERATE_TYPE_DIR} | ||
| # Be compatible with both system-wide but unprivileged env and virtual env where --user is not available | ||
| COMMAND ${PYTHON_EXECUTABLE} -m pip install --user pybind11-stubgen numpy || ${PYTHON_EXECUTABLE} -m pip install pybind11-stubgen numpy | ||
| COMMAND ${PYTHON_EXECUTABLE} -m pybind11_stubgen $<TARGET_FILE_BASE_NAME:python> --output-dir ${WHEEL_TARGET_NAME} | ||
| COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/fix_pyi.py -i ${WHEEL_TARGET_NAME}/$<TARGET_FILE_BASE_NAME:python>.pyi | ||
| WORKING_DIRECTORY ${GENERATE_TYPE_DIR} |
| set(WHEEL_TARGET_NAME "${WHEEL_FILES_DIR}/${PACKAGE_DIR_NAME}") | ||
| set(GENERATE_TYPE_DIR_NAME "onnxruntime_genai-generate-type") | ||
| set(GENERATE_TYPE_DIR "${WHEEL_FILES_DIR}/${GENERATE_TYPE_DIR_NAME}") | ||
| file(MAKE_DIRECTORY ${GENERATE_TYPE_DIR}) |
…ists.txt - fix_pyi.py: use regex with word boundary to strip Oga prefix only at identifier boundaries, and read/write files with explicit UTF-8 encoding - CMakeLists.txt: quote path variables to handle spaces in build directory; install stubgen tools to an isolated target directory to avoid polluting global/user site-packages"
Contributor
|
Have the previously detected issues with pyd type info been resolved in the packaging pipelines? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds pybind11-stubgen based type stub generation for the onnxruntime_genai Python package.
Changes:
Replicates the changes from PR #1742.