32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
"""Make the generated MuJoCo stubs platform-independent and reproducible.
|
|
|
|
``pybind11-stubgen`` introspects the installed ``mujoco`` package, producing two
|
|
kinds of machine-specific output:
|
|
|
|
1. The active OpenGL backend is baked in: ``mujoco.cgl`` on macOS,
|
|
``mujoco.glfw`` on Linux (or ``mujoco.egl`` / ``mujoco.osmesa`` depending on
|
|
``MUJOCO_GL``). mjlab does not type-check against the rendering backend, so we
|
|
drop those stubs and the ``__init__.pyi`` imports referencing them.
|
|
2. Absolute install paths and live object reprs (memory addresses) are embedded
|
|
as default values and ``# value = ...`` comments. These vary by machine and
|
|
run, so we strip them.
|
|
|
|
What remains is a deterministic set of physics-binding stubs that is identical
|
|
on every platform, which lets CI regenerate them and fail if they drift from the
|
|
installed mujoco version.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
STUB_DIR = Path(__file__).parent / "mujoco"
|
|
|
|
# OpenGL context backends mujoco may expose, only one of which exists per
|
|
# platform / MUJOCO_GL setting.
|
|
BACKENDS = ("cgl", "egl", "glfw", "glx", "osmesa")
|
|
|
|
# Backend modules plus the rendering stubs that import them. None are needed to
|
|
# type-check mjlab, and all are platform-specific, so remove them entirely.
|
|
REMOVE = (*BACKENDS, "gl_context.pyi", "renderer.pyi", "rendering")
|
|
|
|
# Backend imports in __init__.pyi, e.g. "from mujoco.cgl import GLContext" and
|
|
# "from . import cgl", along with the names they contribute to __all__.
|
|
DROP_IMPORT = re.compile(
|
|
rf"^(?:from mujoco\.(?:{'|'.join(BACKENDS)}) import GLContext"
|
|
rf"|from \. import (?:{'|'.join(BACKENDS)}))$"
|
|
)
|
|
# Module-level constants in __init__.pyi that exist only on some platforms
|
|
# (macOS exposes is_rosetta / proc_translated for Rosetta detection; Linux does
|
|
# not). Remove them so the stubs match on every platform.
|
|
PLATFORM_ONLY = ("is_rosetta", "proc_translated")
|
|
PLATFORM_ONLY_DEF = re.compile(rf"^(?:{'|'.join(PLATFORM_ONLY)})\b")
|
|
|
|
DROP_FROM_ALL = ("GLContext", *BACKENDS, *PLATFORM_ONLY)
|
|
|
|
# _SYSTEM holds the host platform name ('Darwin' / 'Linux'); drop the value.
|
|
SYSTEM_DEFAULT = re.compile(r"^(_SYSTEM: str) = '\w+'$")
|
|
|
|
# A trailing "# value = ..." comment that captures a live object repr.
|
|
VALUE_COMMENT = re.compile(r" {2,}# value = .*$")
|
|
|
|
# A string constant whose default is an absolute install path.
|
|
PATH_DEFAULT = re.compile(
|
|
r"^(\s*\w+: str) = '[^']*(?:site-packages|/Users/|/home/)[^']*'\s*$"
|
|
)
|
|
|
|
|
|
def remove_backend_stubs() -> None:
|
|
for name in REMOVE:
|
|
target = STUB_DIR / name
|
|
if target.is_dir():
|
|
shutil.rmtree(target)
|
|
elif target.exists():
|
|
target.unlink()
|
|
|
|
|
|
def patch_init() -> None:
|
|
init = STUB_DIR / "__init__.pyi"
|
|
out = []
|
|
for line in init.read_text().splitlines():
|
|
if DROP_IMPORT.match(line.strip()) or PLATFORM_ONLY_DEF.match(line):
|
|
continue
|
|
line = SYSTEM_DEFAULT.sub(r"\1", line)
|
|
if line.startswith("__all__"):
|
|
for name in DROP_FROM_ALL:
|
|
line = line.replace(f"'{name}', ", "")
|
|
out.append(line)
|
|
init.write_text("\n".join(out) + "\n")
|
|
|
|
|
|
def sanitize_machine_specific() -> None:
|
|
for path in STUB_DIR.rglob("*.pyi"):
|
|
out = []
|
|
for line in path.read_text().splitlines():
|
|
if "# value = " in line and ("site-packages" in line or "0x" in line):
|
|
line = VALUE_COMMENT.sub("", line)
|
|
match = PATH_DEFAULT.match(line)
|
|
if match:
|
|
line = match.group(1)
|
|
out.append(line)
|
|
path.write_text("\n".join(out) + "\n")
|
|
|
|
|
|
def main() -> None:
|
|
remove_backend_stubs()
|
|
patch_init()
|
|
sanitize_machine_specific()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|