Install .pyi files

Add rules to generate (as part of the build) and install type stubs for pydrake. Modify wheel build to include these as well. Note that this makes numpy and matplotlib build-time dependencies. (Previously they were only required to test/run pydrake.)
RobotLocomotion · Aug 10, 2022 · 93476a7 · 93476a7
1 parent 3fd412b
commit 93476a7
Show file tree

Hide file tree

Showing 5 changed files with 160 additions and 13 deletions.
diff --git a/bindings/pydrake/BUILD.bazel b/bindings/pydrake/BUILD.bazel
@@ -357,18 +357,6 @@ drake_py_library(
     ],
 )
 
-install(
-    name = "install",
-    install_tests = [
-        ":test/all_install_test.py",
-        ":test/common_install_test.py",
-    ],
-    targets = PY_LIBRARIES + [":all_py"],
-    py_dest = PACKAGE_INFO.py_dest,
-    visibility = ["//visibility:public"],
-    deps = get_drake_py_installs(PY_LIBRARIES_WITH_INSTALL),
-)
-
 drake_cc_googletest(
     name = "documentation_pybind_test",
     deps = [
@@ -692,7 +680,138 @@ drake_py_binary(
     ],
 )
 
-# TODO(mwoehlke-kitware): actually install the .pyi's.
+PYI_FILES = [
+    "pydrake/perception.pyi",
+    "pydrake/autodiffutils.pyi",
+    "pydrake/math.pyi",
+    "pydrake/trajectories.pyi",
+    "pydrake/visualization/all.pyi",
+    "pydrake/visualization/__init__.pyi",
+    "pydrake/visualization/meldis.pyi",
+    "pydrake/visualization/plotting.pyi",
+    "pydrake/all.pyi",
+    "pydrake/lcm.pyi",
+    "pydrake/manipulation/all.pyi",
+    "pydrake/manipulation/schunk_wsg.pyi",
+    "pydrake/manipulation/__init__.pyi",
+    "pydrake/manipulation/planner.pyi",
+    "pydrake/manipulation/kuka_iiwa.pyi",
+    "pydrake/forwarddiff.pyi",
+    "pydrake/multibody/math.pyi",
+    "pydrake/multibody/all.pyi",
+    "pydrake/multibody/meshcat.pyi",
+    "pydrake/multibody/jupyter_widgets.pyi",
+    "pydrake/multibody/plant.pyi",
+    "pydrake/multibody/tree.pyi",
+    "pydrake/multibody/optimization.pyi",
+    "pydrake/multibody/parsing.pyi",
+    "pydrake/multibody/benchmarks.pyi",
+    "pydrake/multibody/__init__.pyi",
+    "pydrake/multibody/inverse_kinematics.pyi",
+    "pydrake/geometry.pyi",
+    "pydrake/polynomial.pyi",
+    "pydrake/__init__.pyi",
+    "pydrake/examples/pendulum.pyi",
+    "pydrake/examples/multibody/__init__.pyi",
+    "pydrake/examples/multibody/run_planar_scenegraph_visualizer.pyi",
+    "pydrake/examples/multibody/cart_pole_passive_simulation.pyi",
+    "pydrake/examples/multibody/pendulum_lqr_monte_carlo_analysis.pyi",
+    "pydrake/examples/van_der_pol.pyi",
+    "pydrake/examples/manipulation_station.pyi",
+    "pydrake/examples/_examples_extra.pyi",
+    "pydrake/examples/__init__.pyi",
+    "pydrake/examples/rimless_wheel.pyi",
+    "pydrake/examples/compass_gait.pyi",
+    "pydrake/examples/quadrotor.pyi",
+    "pydrake/examples/acrobot.pyi",
+    "pydrake/systems/perception.pyi",
+    "pydrake/systems/trajectory_optimization.pyi",
+    "pydrake/systems/drawing.pyi",
+    "pydrake/systems/all.pyi",
+    "pydrake/systems/lcm.pyi",
+    "pydrake/systems/jupyter_widgets.pyi",
+    "pydrake/systems/rendering.pyi",
+    "pydrake/systems/framework.pyi",
+    "pydrake/systems/analysis.pyi",
+    "pydrake/systems/pyplot_visualizer.pyi",
+    "pydrake/systems/planar_scenegraph_visualizer.pyi",
+    "pydrake/systems/primitives.pyi",
+    "pydrake/systems/__init__.pyi",
+    "pydrake/systems/scalar_conversion.pyi",
+    "pydrake/systems/_resample_interp1d.pyi",
+    "pydrake/systems/controllers.pyi",
+    "pydrake/systems/sensors.pyi",
+    "pydrake/systems/meshcat_visualizer.pyi",
+    "pydrake/common/cpp_template.pyi",
+    "pydrake/common/value.pyi",
+    "pydrake/common/yaml.pyi",
+    "pydrake/common/all.pyi",
+    "pydrake/common/eigen_geometry.pyi",
+    "pydrake/common/_module_py.pyi",
+    "pydrake/common/schema.pyi",
+    "pydrake/common/__init__.pyi",
+    "pydrake/common/cpp_param.pyi",
+    "pydrake/common/jupyter.pyi",
+    "pydrake/common/deprecation.pyi",
+    "pydrake/common/compatibility.pyi",
+    "pydrake/common/pybind11_version.pyi",
+    "pydrake/common/containers.pyi",
+    "pydrake/tutorials.pyi",
+    "pydrake/stubgen.pyi",
+    "pydrake/symbolic.pyi",
+    "pydrake/solvers/ipopt.pyi",
+    "pydrake/solvers/osqp.pyi",
+    "pydrake/solvers/scs.pyi",
+    "pydrake/solvers/all.pyi",
+    "pydrake/solvers/_extra.pyi",
+    "pydrake/solvers/clp.pyi",
+    "pydrake/solvers/ibex.pyi",
+    "pydrake/solvers/branch_and_bound.pyi",
+    "pydrake/solvers/mixed_integer_rotation_constraint.pyi",
+    "pydrake/solvers/csdp.pyi",
+    "pydrake/solvers/mixed_integer_optimization_util.pyi",
+    "pydrake/solvers/nlopt.pyi",
+    "pydrake/solvers/dreal.pyi",
+    "pydrake/solvers/sdpa_free_format.pyi",
+    "pydrake/solvers/augmented_lagrangian.pyi",
+    "pydrake/solvers/__init__.pyi",
+    "pydrake/solvers/mosek.pyi",
+    "pydrake/solvers/mathematicalprogram.pyi",
+    "pydrake/solvers/snopt.pyi",
+    "pydrake/solvers/gurobi.pyi",
+]
+
+genrule(
+    name = "pydrake_pyi",
+    srcs = [],
+    outs = PYI_FILES,
+    cmd = "$(location :stubgen) --package=pydrake --output=$(RULEDIR)",
+    tools = [":stubgen"],
+)
+
+TYPED_FILES = ["pydrake/py.typed"]
+
+genrule(
+    name = "pydrake_typed",
+    srcs = [],
+    outs = TYPED_FILES,
+    cmd = "touch $@",
+    tools = [":stubgen"],
+)
+
+install(
+    name = "install",
+    install_tests = [
+        ":test/all_install_test.py",
+        ":test/common_install_test.py",
+    ],
+    targets = PY_LIBRARIES + [":all_py"],
+    py_dest = PACKAGE_INFO.py_dest,
+    data = PYI_FILES + TYPED_FILES,
+    data_dest = "@PYTHON_SITE_PACKAGES@",
+    visibility = ["//visibility:public"],
+    deps = get_drake_py_installs(PY_LIBRARIES_WITH_INSTALL),
+)
 
 drake_py_unittest(
     name = "stubgen_test",

diff --git a/bindings/pydrake/test/all_install_test.py b/bindings/pydrake/test/all_install_test.py
@@ -25,6 +25,17 @@ def test_install(self):
             env=tool_env
         )
 
+    def test_type_stubs(self):
+        # Get pydrake directory.
+        install_dir = install_test_helper.get_install_dir()
+        site_packages = \
+            install_test_helper.get_python_site_packages_dir(install_dir)
+        # Ensure some files are present.
+        self.assertTrue(
+            os.path.exists(os.path.join(site_packages, 'pydrake', 'py.typed')))
+        self.assertTrue(
+            os.path.exists(os.path.join(site_packages, 'pydrake', 'all.pyi')))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tools/wheel/image/provision-python.sh b/tools/wheel/image/provision-python.sh
@@ -16,6 +16,8 @@ ${PYTHON} -m venv /usr/local
 
 pip install \
     lxml \
+    matplotlib \
+    numpy \
     pyyaml \
     semantic-version \
     setuptools \

diff --git a/tools/wheel/image/setup.py b/tools/wheel/image/setup.py
@@ -94,6 +94,8 @@ def _actually_find_packages():
       include_package_data=True,
       package_data={
           '': find_data_files(
+              'pydrake/py.typed',
+              'pydrake/**/*.pyi',
               'pydrake/**/*.so',
               'pydrake/lib/**',
               'pydrake/doc/**',

diff --git a/tools/wheel/test/test-wheel.sh b/tools/wheel/test/test-wheel.sh
@@ -56,3 +56,16 @@ pypi_wheel_max_size = 100 << 20  # 100 MiB
 fail_message = f"Wheel is too large ({wheel_size_in_mib:.2f} MiB) for PyPI"
 assert wheel_size < pypi_wheel_max_size, fail_message
 EOF
+
+python - "$1" << EOF
+import os
+import pydrake
+
+def assert_exists(path):
+    assert os.path.exists(path), f"{path!r} does not exist!"
+
+# Check that type information files are present.
+pydrake_dir = pydrake.__path__[0]
+assert_exists(os.path.join(pydrake_dir, 'py.typed'))
+assert_exists(os.path.join(pydrake_dir, 'all.pyi'))
+EOF