Python Types are missing :(

[Artur-Galstyan]

Is your feature request related to a problem? Please describe.
Pyright LSP cannot find the MjData, MjModel etc. types 😢

Describe the solution you'd like
A potential solution could be adding python stubs, like amujoco.pyi file or something?

FYI: I was checking out Brax when I noticed this

[kevinzakka]

Hi @Artur-Galstyan, I was able to use pybind11-stubgen and point Visual Code to the generated stubs to get type hints.

pip install pybind11-stubgen
pybind11-stubgen mujoco -o ~/typings/  # Installs mujoco stubs dir in $HOME/typings/

Now in visual code, I added the following to my settings.json:

"python.languageServer": "Pylance",
"python.analysis.stubPath": "~/typings",

and the result:

It's not perfect but I think it generates hints for a good chunk of the API. Hope that helps!

[mluogh-xdof]

I was able to get this to work using Pylance typechecking in CI as well.

I had to patch the _enum.pyi file created because the function bindings created expect an int type which works in C because of the way enums are basically ints there, but the Python typechecker throws an error. So we have to patch them to be in IntEnums.

Here's the script I used for this and I put it in the typings folder in my project: https://gist.github.com/mluogh-xdof/805cd47daf7d3690495c74cb1141ddb7

I also ran it to generate bindings using Annotated because there were some typing errors with the standard numpy annotations due to dimension annotations (which might alternatively be fixed by this currently unmerged PR pybind/pybind11-stubgen#264)

Now it doesn't just give typehints but gives reasonable typechecking and can be shared across teams/CI/etc.

tldr workflow:

In the root dir:

pybind11-stubgen mujoco -o typings/  --numpy-array-wrap-with-annotated

Then:

python typings/patch_mujoco_enums.py typings/mujoco/_enums.pyi

Then in pyproject.toml:

[tool.pyright]
stubPath = "typings"

EDIT: actually seems like a recent enough version mujoco is compiled with a version of pybind11 that will use typing.SupportsInt and make this issue go away automatically: pybind/pybind11#5540?

so this will work:

pybind11-stubgen mujoco -o ~/typings

[kevinzakka]

Thanks @mluogh-xdof (and hi @wuphilipp). FYI @saran-t