stubgen improvements

python/py_package/__init__.pyi

@@ -34,4 +34,4 @@ from . import utils
 from . import version
 from . import wrapper
 __all__ = ['ActorBuilder', 'ArticulationBuilder', 'Component', 'CudaArray', 'Engine', 'Entity', 'Path', 'PinocchioModel', 'Pose', 'SapienRenderer', 'Scene', 'SceneConfig', 'System', 'Widget', 'asset', 'internal_renderer', 'math', 'os', 'physx', 'pkg_resources', 'platform', 'profile', 'pysapien', 'pysapien_pinocchio', 'render', 'render_server', 'set_log_level', 'simsense', 'utils', 'version', 'warn', 'wrapper']
-__version__: str = '3.0.0.dev20240127+33ca7c50'
+__version__: str = '3.0.0.dev20240127+35c8314e'

python/py_package/pysapien/__init__.pyi

@@ -1,13 +1,15 @@
 from __future__ import annotations
+import jax
 import numpy
+import torch
 import typing
 from . import internal_renderer
 from . import math
 from . import physx
 from . import render
 from . import simsense
 __all__ = ['Component', 'CudaArray', 'Entity', 'Pose', 'Profiler', 'Scene', 'System', 'internal_renderer', 'math', 'physx', 'profile', 'render', 'set_log_level', 'simsense']
-_T = typing.TypeVar("_T", bound=Component)
+_T = typing.TypeVar("_T", bound=sapien.Component)
 class Component:
     entity_pose: Pose
     name: str
@@ -47,9 +49,9 @@ class CudaArray:
         ...
     def dlpack(self) -> typing.Any:
         ...
-    def jax(self) -> typing.Any:
+    def jax(self) -> jax.Array:
         ...
-    def torch(self) -> typing.Any:
+    def torch(self) -> torch.Tensor:
         ...
     @property
     def __cuda_array_interface__(self) -> dict:
@@ -151,7 +153,7 @@ class Pose:
     def to_transformation_matrix(self) -> numpy.ndarray[tuple[typing.Literal[4], typing.Literal[4]], numpy.dtype[numpy.float32]]:
         ...
 class Profiler:
-    def __call__(self, arg0: typing.Callable) -> cpp_function:
+    def __call__(self, arg0: typing.Callable) -> typing.Callable:
         ...
     def __enter__(self) -> None:
         ...
@@ -203,7 +205,7 @@ class System:
 def profile(arg0: str) -> Profiler:
     ...
 @typing.overload
-def profile(arg0: typing.Callable) -> cpp_function:
+def profile(arg0: typing.Callable) -> typing.Callable:
     ...
 def set_log_level(level: str) -> None:
     ...

python/stubgen.py

@@ -33,6 +33,45 @@
 # pybind11-stubgen == 2.4.2
 
 
+class FixTorchJax(IParser):
+    def handle_class_member(
+        self, path: QualifiedName, class_: type, obj: Any
+    ) -> Docstring | Alias | Class | list[Method] | Field | Property | None:
+        result = super().handle_class_member(path, class_, obj)
+        if not isinstance(result, list):
+            return result
+
+        for method in result:
+            if str(method.function.name) == "torch":
+                method.function.returns = ResolvedType(
+                    QualifiedName.from_str("torch.Tensor")
+                )
+
+        for method in result:
+            if str(method.function.name) == "jax":
+                method.function.returns = ResolvedType(
+                    QualifiedName.from_str("jax.Array")
+                )
+
+        return result
+
+    def handle_module(
+        self, path: QualifiedName, module: types.ModuleType
+    ) -> Module | None:
+        result = super().handle_module(path, module)
+        if result is None:
+            return None
+        if str(path).endswith("pysapien"):
+            result.imports.add(
+                Import(name=None, origin=QualifiedName.from_str("torch")),
+            )
+            result.imports.add(
+                Import(name=None, origin=QualifiedName.from_str("jax")),
+            )
+
+        return result
+
+
 class FixCppFunction(IParser):
     def handle_function(self, path: QualifiedName, func: Any) -> list[Function]:
         result = super().handle_function(path, func)
@@ -146,7 +185,7 @@ def handle_module(
             result.type_vars.append(
                 TypeVar_(
                     name=Identifier("_T"),
-                    bound=ResolvedType(QualifiedName.from_str("Component")),
+                    constraints=[ResolvedType((QualifiedName.from_str("Component")))],
                 ),
             )
 
@@ -171,6 +210,7 @@ class Parser(
         FixMissingImports,
         FilterTypingModuleAttributes,
         FixPEP585CollectionNames,
+        FixTorchJax,
         FixCppFunction,
         FixNoneParameterType,
         FixTypingTypeNames,