Automatically generate pyi stubs for built extensions

- Fixes #1

Author: virtuald

examples/demo/.gitignore

@@ -1,8 +1,9 @@
-*.py[cod]
+*.py[ciod]
 *.so
 *.dylib
 *.egg-info
 *.whl
+py.typed
 
 /build
 /dist

robotpy_build/command/build_ext.py

@@ -93,14 +93,17 @@ def build_extensions(self):
             install_root = get_install_root(self)
 
             for ext in self.extensions:
-                relink_extension(
+                libs = relink_extension(
                     install_root,
                     self.get_ext_fullpath(ext.name),
                     self.get_ext_filename(ext.name),
                     ext.rpybuild_wrapper,
                     self.rpybuild_pkgcfg,
                 )
 
+                # Used in build_pyi
+                ext.rpybuild_libs = libs
+
     def run(self):
 
         # files need to be generated before building can occur
@@ -111,6 +114,9 @@ def run(self):
 
         build_ext.run(self)
 
+        # pyi can only be built after ext is built
+        self.run_command("build_pyi")
+
     def get_libraries(self, ext):
         libraries = build_ext.get_libraries(self, ext)
 

robotpy_build/command/build_pyi.py

@@ -0,0 +1,162 @@
+import importlib.util
+import json
+import os
+from os.path import abspath, exists, dirname, join
+import subprocess
+import sys
+
+import pybind11_stubgen
+from setuptools import Command
+from distutils.errors import DistutilsError
+
+from .util import get_install_root
+
+
+class GeneratePyiError(DistutilsError):
+    pass
+
+
+class BuildPyi(Command):
+
+    command_name = "build_pyi"
+    description = "Generates pyi files from built extensions"
+
+    user_options = [("build-lib=", "d", 'directory to "build" (copy) to')]
+
+    def initialize_options(self):
+        self.build_lib = None
+
+    def finalize_options(self):
+        self.set_undefined_options("build", ("build_lib", "build_lib"))
+
+    def run(self):
+        # cannot build pyi files when cross-compiling
+        if (
+            "_PYTHON_HOST_PLATFORM" in os.environ
+            or os.environ.get("RPYBUILD_SKIP_PYI") == "1"
+        ):
+            return
+
+        # Gather information for needed stubs
+        data = {"mapping": {}, "stubs": []}
+
+        # OSX-specific: need to set DYLD_LIBRARY_PATH otherwise modules don't
+        # work. Luckily, that information was computed when building the
+        # extensions...
+        env = os.environ.copy()
+        dyld_path = set()
+
+        # Requires information from build_ext to work
+        build_ext = self.distribution.get_command_obj("build_ext")
+        if build_ext.inplace:
+            data["out"] = get_install_root(self)
+        else:
+            data["out"] = self.build_lib
+
+        # Ensure that the associated packages can always be found locally
+        for wrapper in build_ext.wrappers:
+            pkgdir = wrapper.package_name.split(".")
+            init_py = abspath(join(self.build_lib, *pkgdir, "__init__.py"))
+            if exists(init_py):
+                data["mapping"][wrapper.package_name] = init_py
+
+        # Ensure that the built extension can always be found
+        for ext in build_ext.extensions:
+            fname = build_ext.get_ext_filename(ext.name)
+            data["mapping"][ext.name] = abspath(join(self.build_lib, fname))
+            data["stubs"].append(ext.name)
+
+            rpybuild_libs = getattr(ext, "rpybuild_libs", None)
+            if rpybuild_libs:
+                for pth, _ in rpybuild_libs.values():
+                    dyld_path.add(dirname(pth))
+
+        # Don't do anything if nothing is needed
+        if not data["stubs"]:
+            return
+
+        # OSX-specific
+        if dyld_path:
+            dyld_path = ":".join(dyld_path)
+            if "DYLD_LIBRARY_PATH" in env:
+                dyld_path += ":" + env["DYLD_LIBRARY_PATH"]
+            env["DYLD_LIBRARY_PATH"] = dyld_path
+
+        data_json = json.dumps(data)
+
+        # Execute in a subprocess in case it crashes
+        args = [sys.executable, "-m", __name__]
+        try:
+            subprocess.run(args, input=data_json.encode("utf-8"), env=env, check=True)
+        except subprocess.CalledProcessError:
+            raise GeneratePyiError(
+                "Failed to generate .pyi file (see above, or set RPYBUILD_SKIP_PYI=1 to ignore) via %s"
+                % (args,)
+            ) from None
+
+
+class _PackageFinder:
+    """
+    Custom loader to allow loading built modules from their location
+    in the build directory (as opposed to their install location)
+    """
+
+    mapping = {}
+
+    @classmethod
+    def find_spec(cls, fullname, path, target=None):
+        m = cls.mapping.get(fullname)
+        if m:
+            return importlib.util.spec_from_file_location(fullname, m)
+
+
+def generate_pyi(module_name: str, pyi_filename: str):
+
+    print("generating", pyi_filename)
+
+    pybind11_stubgen.FunctionSignature.n_invalid_signatures = 0
+    module = pybind11_stubgen.ModuleStubsGenerator(module_name)
+    module.parse()
+    if pybind11_stubgen.FunctionSignature.n_invalid_signatures > 0:
+        print("FAILED to generate pyi for", module_name, file=sys.stderr)
+        return False
+
+    module.write_setup_py = False
+    with open(pyi_filename, "w") as fp:
+        fp.write("#\n# AUTOMATICALLY GENERATED FILE, DO NOT EDIT!\n#\n\n")
+        fp.write("\n".join(module.to_lines()))
+
+    typed = join(dirname(pyi_filename), "py.typed")
+    print("generating", typed)
+    if not exists(typed):
+        with open(typed, "w") as fp:
+            pass
+
+    return True
+
+
+def main():
+
+    cfg = json.load(sys.stdin)
+
+    # Configure custom loader
+    _PackageFinder.mapping = cfg["mapping"]
+    sys.meta_path.insert(0, _PackageFinder)
+
+    # Generate pyi modules
+    sys.argv = [
+        "<dummy>",
+        "--no-setup-py",
+        "--log-level=WARNING",
+        "--root-module-suffix=",
+        "--ignore-invalid",
+        "defaultarg",
+        "-o",
+        cfg["out"],
+    ] + cfg["stubs"]
+
+    pybind11_stubgen.main()
+
+
+if __name__ == "__main__":
+    main()

robotpy_build/relink_libs.py

@@ -137,7 +137,7 @@ def relink_extension(
     extension_rel: str,
     pkg: PkgCfg,
     pkgcfg: PkgCfgProvider,
-):
+) -> LibsDict:
     """
     Given an extension, relink it
 
@@ -158,3 +158,4 @@ def relink_extension(
         )
     }
     _fix_libs(to_fix, libs)
+    return libs

robotpy_build/setup.py

@@ -22,6 +22,7 @@ def finalize_options(self):
 from .command.build_dl import BuildDl
 from .command.build_gen import BuildGen
 from .command.build_ext import BuildExt
+from .command.build_pyi import BuildPyi
 from .command.develop import Develop
 
 from .overrides import apply_overrides
@@ -112,6 +113,7 @@ def prepare(self):
             "build_dl": BuildDl,
             "build_gen": BuildGen,
             "build_ext": BuildExt,
+            "build_pyi": BuildPyi,
             "develop": Develop,
         }
         if bdist_wheel:

setup.cfg

@@ -32,6 +32,7 @@ install_requires =
     toposort
     pyyaml >= 5.1
     patch == 1.*
+    pybind11-stubgen >= 0.8.1
     dataclasses; python_version < '3.7'
     delocate; platform_system == 'Darwin'
 setup_requires =

tests/cpp/.gitignore

@@ -1,10 +1,12 @@
 
-*.pyc
+*.py[ciod]
 
 *.so
 *.dll
 *.pyd
 *.dylib
+py.typed
+
 
 /build
 /dist

tests/cpp/rpytest/ft/include/inheritance/ibase.h

@@ -52,9 +52,9 @@ struct IBase
         return 7;
     }
 
-    virtual void protectedOutMethod(int *out, int in)
+    virtual void protectedOutMethod(int *out, int inp)
     {
-        *out = in + 5;
+        *out = inp + 5;
     }
 };