cmake(python): Add GENERATE_PYTHON_STUB option (Fix #3944)

2023-10-16 15:06:26 +02:00
parent f36ad15224
commit e1f9febd8e
8 changed files with 77 additions and 7 deletions
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -362,6 +362,9 @@ if(BUILD_PYTHON)
    "NOT BUILD_DEPS" ON)
  message(STATUS "Python: Build pybind11_protobuf: ${BUILD_pybind11_protobuf}")

+  option(GENERATE_PYTHON_STUB "Generate Python stub file (.pyi)" ON)
+  message(STATUS "Python: Generate stub file: ${GENERATE_PYTHON_STUB}")
+
 CMAKE_DEPENDENT_OPTION(BUILD_PYTHON_DOC "Build the Python doc" OFF "NOT BUILD_DOC" ON)
  message(STATUS "Python: Build doc: ${BUILD_PYTHON_DOC}")

--- a/bazel/notebook_requirements.in
+++ b/bazel/notebook_requirements.in
@@ -5,6 +5,7 @@ protobuf==4.24.4
 scipy==1.11.1

 # OR-Tools build dependencies
+mypy==1.6.1
 mypy-protobuf==3.5.0
 virtualenv==20.24.2
 black==23.7.0
--- a/bazel/notebook_requirements.txt
+++ b/bazel/notebook_requirements.txt
@@ -158,8 +158,12 @@ matplotlib-inline==0.1.6
    #   ipython
 mistune==3.0.1
    # via nbconvert
+mypy==1.6.1
+    # via -r bazel/notebook_requirements.in
 mypy-extensions==1.0.0
-    # via black
+    # via
+    #   black
+    #   mypy
 mypy-protobuf==3.5.0
    # via -r bazel/notebook_requirements.in
 nbclient==0.8.0
@@ -340,6 +344,8 @@ traitlets==5.9.0
    #   qtconsole
 types-protobuf==4.24.0.0
    # via mypy-protobuf
+typing-extensions==4.8.0
+    # via mypy
 tzdata==2023.3
    # via pandas
 uri-template==1.3.0
--- a/bazel/ortools_requirements.in
+++ b/bazel/ortools_requirements.in
@@ -5,6 +5,7 @@ protobuf==4.24.4
 scipy==1.11.1

 # OR-Tools build dependencies
+mypy==1.6.1
 mypy-protobuf==3.5.0
 virtualenv==20.24.2
 black==23.7.0
--- a/bazel/ortools_requirements.txt
+++ b/bazel/ortools_requirements.txt
@@ -14,8 +14,12 @@ distlib==0.3.7
    # via virtualenv
 filelock==3.12.2
    # via virtualenv
+mypy==1.6.1
+    # via -r bazel/ortools_requirements.in
 mypy-extensions==1.0.0
-    # via black
+    # via
+    #   black
+    #   mypy
 mypy-protobuf==3.5.0
    # via -r bazel/ortools_requirements.in
 numpy==1.25.2
@@ -47,6 +51,8 @@ six==1.16.0
    # via python-dateutil
 types-protobuf==4.24.0.0
    # via mypy-protobuf
+typing-extensions==4.8.0
+    # via mypy
 tzdata==2023.3
    # via pandas
 virtualenv==20.24.2
--- a/cmake/README.md
+++ b/cmake/README.md
@@ -253,6 +253,8 @@ cmake -S. -Bbuild -LH
 | `BUILD_FAT_JAR`  | OFF | Build a `ortools-java` .jar that includes all of its own Maven dependencies, including the native package<br>Only available if `BUILD_JAVA=ON` |
 | | | |
 | `BUILD_pybind11` | `BUILD_DEPS` | Static build the pybind11 libraries<br>**Forced** to ON if `BUILD_DEPS=ON`<br>Only available if `BUILD_PYTHON=ON` |
+| `BUILD_pybind11_protobuf` | `BUILD_DEPS` | Static build the pybind11_protobuf libraries<br>**Forced** to ON if `BUILD_DEPS=ON`<br>Only available if `BUILD_PYTHON=ON` |
+| `GENERATE_PYTHON_STUB` | ON | Generate python stub files<br>Only available if `BUILD_PYTHON=ON` |
 | `BUILD_VENV` | `BUILD_TESTING` | Create python venv in `BINARY_DIR/python/venv`<br>**Forced** to ON if `BUILD_TESTING=ON`<br>Only available if `BUILD_PYTHON=ON` |
 | `VENV_USE_SYSTEM_SITE_PACKAGES` | OFF | Python venv can use system site package (e.g. `py3-numpy` on Alpine)<br>Only available if `BUILD_PYTHON=ON` and `BUILD_VENV=ON` |
 | `FETCH_PYTHON_DEPS`  | `BUILD_DEPS` | Fetch python modules needed to build ortools package<br>Only available if `BUILD_PYTHON=ON` |
--- a/cmake/python.cmake
+++ b/cmake/python.cmake
@@ -308,6 +308,56 @@ configure_file(
  ${PROJECT_BINARY_DIR}/python/README.txt
  COPYONLY)

+# Generate Stub
+if(GENERATE_PYTHON_STUB)
+# Look for required python modules
+search_python_module(
+  NAME mypy
+  PACKAGE mypy
+  NO_VERSION)
+
+find_program(
+  stubgen_EXECUTABLE
+  NAMES stubgen stubgen.exe
+  REQUIRED
+)
+
+add_custom_command(
+  OUTPUT python/stub/timestamp
+  COMMAND ${CMAKE_COMMAND} -E remove_directory stub
+  COMMAND ${CMAKE_COMMAND} -E make_directory stub
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.init.python.init --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.algorithms.python.knapsack_solver --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.graph.python.linear_sum_assignment --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.graph.python.max_flow --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.graph.python.min_cost_flow --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.constraint_solver.pywrapcp --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.linear_solver.pywraplp --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.linear_solver.python.model_builder_helper --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.pdlp.python.pdlp --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.sat.python.swig_helper --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.scheduling.python.rcpsp --output .
+  COMMAND ${stubgen_EXECUTABLE} -p ortools.util.python.sorted_interval_list --output .
+  COMMAND ${CMAKE_COMMAND} -E touch ${PROJECT_BINARY_DIR}/python/stub/timestamp
+  MAIN_DEPENDENCY
+    ortools/python/setup.py.in
+  DEPENDS
+    init_pybind11
+    knapsack_solver_pybind11
+    linear_sum_assignment_pybind11
+    max_flow_pybind11
+    min_cost_flow_pybind11
+    pywrapcp
+    pywraplp
+    model_builder_helper_pybind11
+    pdlp_pybind11
+    swig_helper_pybind11
+    rcpsp_pybind11
+    sorted_interval_list_pybind11
+  WORKING_DIRECTORY python
+  COMMAND_EXPAND_LISTS)
+endif()
+
 # Look for required python modules
 search_python_module(
  NAME setuptools
@@ -357,6 +407,7 @@ add_custom_command(
    swig_helper_pybind11
    rcpsp_pybind11
    sorted_interval_list_pybind11
+    $<$<BOOL:${GENERATE_PYTHON_STUB}>:python/stub/timestamp>
  BYPRODUCTS
    python/${PYTHON_PROJECT}
    python/${PYTHON_PROJECT}.egg-info
--- a/ortools/python/setup.py.in
+++ b/ortools/python/setup.py.in
@@ -50,21 +50,21 @@ setup(
    ],
    package_data={
        '@PYTHON_PROJECT@':[$<$<STREQUAL:$<TARGET_PROPERTY:@PROJECT_NAME@,TYPE>,SHARED_LIBRARY>:'.libs/*','../$<TARGET_SONAME_FILE_NAME:@PROJECT_NAME@>'>],
-        '@PYTHON_PROJECT@.init.python':['$<TARGET_FILE_NAME:init_pybind11>'],
-        '@PYTHON_PROJECT@.algorithms.python':['$<TARGET_FILE_NAME:knapsack_solver_pybind11>'],
+        '@PYTHON_PROJECT@.init.python':['$<TARGET_FILE_NAME:init_pybind11>', '*.pyi'],
+        '@PYTHON_PROJECT@.algorithms.python':['$<TARGET_FILE_NAME:knapsack_solver_pybind11>', '*.pyi'],
        '@PYTHON_PROJECT@.bop':['*.pyi'],
        '@PYTHON_PROJECT@.glop':['*.pyi'],
        '@PYTHON_PROJECT@.graph.python':[
            '$<TARGET_FILE_NAME:linear_sum_assignment_pybind11>',
            '$<TARGET_FILE_NAME:max_flow_pybind11>',
-            '$<TARGET_FILE_NAME:min_cost_flow_pybind11>'
-        ],
+            '$<TARGET_FILE_NAME:min_cost_flow_pybind11>',
+            '*.pyi'],
        '@PYTHON_PROJECT@.constraint_solver':['$<TARGET_FILE_NAME:pywrapcp>', '*.pyi'],
        '@PYTHON_PROJECT@.linear_solver':['$<TARGET_FILE_NAME:pywraplp>', '*.pyi'],
        '@PYTHON_PROJECT@.linear_solver.python':['$<TARGET_FILE_NAME:model_builder_helper_pybind11>', '*.pyi'],
        '@PYTHON_PROJECT@.packing':['*.pyi'],
        '@PYTHON_PROJECT@.pdlp':['*.pyi'],
-        '@PYTHON_PROJECT@.pdlp.python':['$<TARGET_FILE_NAME:pdlp_pybind11>'],
+        '@PYTHON_PROJECT@.pdlp.python':['$<TARGET_FILE_NAME:pdlp_pybind11>', '*.pyi'],
        '@PYTHON_PROJECT@.sat':['*.pyi'],
        '@PYTHON_PROJECT@.sat.colab':['*.pyi'],
        '@PYTHON_PROJECT@.sat.python':['$<TARGET_FILE_NAME:swig_helper_pybind11>', '*.pyi'],