Updated pybind11-stubgen and fixed annotation issues

Author: Tim Schneider

custom_stubgen.py

@@ -1,25 +1,63 @@
 #!/usr/bin/env python3
-import pybind11_stubgen
-import re
+from typing import Dict
 
+from pybind11_stubgen import *
+from pybind11_stubgen.structs import *
 
-def add_union(alternative_type: str):
-    def inner(match, _alternative_type=alternative_type):
-        return "typing.Union[{}, {}]".format(match.group(0), _alternative_type)
 
-    return inner
+class CustomWriter(Writer):
+    def __init__(self, implicit_conversions: Dict[str, str], stub_ext: str = "pyi"):
+        super().__init__(stub_ext=stub_ext)
+        self.implicit_conversions = {
+            QualifiedName.from_str(k): QualifiedName.from_str(v) for k, v in implicit_conversions.items()
+        }
 
+    def _patch_function(self, function: Function):
+        for argument in function.args:
+            if argument.annotation is not None and argument.annotation.name in self.implicit_conversions:
+                converted_type = ResolvedType(self.implicit_conversions[argument.annotation.name])
+                argument.annotation = ResolvedType(
+                    QualifiedName.from_str("typing.Union"), [argument.annotation, converted_type])
 
-if __name__ == '__main__':
-    implicit_conversions = {
-        "bool": "Condition",
-        "float": "RelativeDynamicsFactor",
-        "Affine": "RobotPose",
-    }
+    def write_module(self, module: Module, printer: Printer, to: Path, sub_dir: Optional[Path] = None):
+        for cls in module.classes:
+            for method in cls.methods:
+                self._patch_function(method.function)
+            for prop in cls.properties:
+                if prop.setter is not None:
+                    self._patch_function(prop.setter)
+        super().write_module(module, printer, to, sub_dir=sub_dir)
 
-    pybind11_stubgen.StubsGenerator.GLOBAL_CLASSNAME_REPLACEMENTS.update({
-        re.compile("({})".format(orig_type)): add_union(alt_type)
-        for alt_type, orig_type in implicit_conversions.items()
-    })
 
-    pybind11_stubgen.main()
+IMPLICIT_CONVERSIONS = {
+    "bool": "Condition",
+    "float": "RelativeDynamicsFactor",
+    "Affine": "RobotPose",
+}
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(name)s - [%(levelname)7s] %(message)s",
+    )
+    args = arg_parser().parse_args()
+
+    parser = stub_parser_from_args(args)
+
+    printer = Printer(invalid_expr_as_ellipses=not args.print_invalid_expressions_as_is)
+
+    out_dir, sub_dir = to_output_and_subdir(
+        output_dir=args.output_dir,
+        module_name=args.module_name,
+        root_suffix=args.root_suffix,
+    )
+
+    run(
+        parser,
+        printer,
+        args.module_name,
+        out_dir,
+        sub_dir=sub_dir,
+        dry_run=args.dry_run,
+        writer=CustomWriter(IMPLICIT_CONVERSIONS, stub_ext=args.stub_extension),
+    )

generate_stubs.bash

@@ -18,9 +18,7 @@ ${PYTHON} -m venv "${STUBS_GEN_DIR}/venv"
 source "${STUBS_GEN_DIR}/venv/bin/activate"
 (cd ${SCRIPT_DIR} && python "setup.py" egg_info -e "${STUBS_GEN_DIR}" > /dev/null)
 pip install -r "${STUBS_GEN_DIR}/franky_panda.egg-info/requires.txt" > /dev/null
-pip install pybind11-stubgen==0.16.0 > /dev/null
+pip install pybind11-stubgen==2.2.2 > /dev/null
 
-PYTHONPATH="${LIB_DIR}" "${SCRIPT_DIR}/custom_stubgen.py" --no-setup-py -o "${STUBS_GEN_DIR}" _franky
-# TODO: remove this workaround once pybind11-stubgen is updated
-sed "s/import _franky/import franky._franky as _franky/" "${STUBS_GEN_DIR}/_franky-stubs/__init__.pyi" > "${LIB_DIR}/_franky.pyi"
+PYTHONPATH="${LIB_DIR}" "${SCRIPT_DIR}/custom_stubgen.py" -o "${LIB_DIR}" _franky