Stricter pyright Setup

ambersim/utils/_internal_utils.py

@@ -4,7 +4,7 @@
 from ambersim import ROOT
 
 
-def _check_filepath(filepath: Union[str, Path]) -> str:
+def check_filepath(filepath: Union[str, Path]) -> str:
     """Checks validity of a filepath for model loading."""
     assert isinstance(filepath, (str, Path))
 
@@ -19,7 +19,7 @@ def _check_filepath(filepath: Union[str, Path]) -> str:
     return filepath
 
 
-def _rmtree(f: Path):
+def rmtree(f: Path):
     """Recursively deletes a directory using pathlib.
 
     See: https://stackoverflow.com/a/66552066
@@ -28,5 +28,5 @@ def _rmtree(f: Path):
         f.unlink()
     else:
         for child in f.iterdir():
-            _rmtree(child)
+            rmtree(child)
         f.rmdir()

ambersim/utils/conversion_utils.py

@@ -5,7 +5,7 @@
 import mujoco as mj
 import trimesh
 
-from ambersim.utils._internal_utils import _check_filepath
+from ambersim.utils._internal_utils import check_filepath
 
 
 def save_model_xml(filepath: Union[str, Path], output_name: Optional[str] = None) -> None:
@@ -21,7 +21,7 @@ def save_model_xml(filepath: Union[str, Path], output_name: Optional[str] = None
     """
     try:
         # loading model and saving XML
-        filepath = _check_filepath(filepath)
+        filepath = check_filepath(filepath)
         _model = mj.MjModel.from_xml_path(filepath)
         if output_name is None:
             output_name = filepath.split("/")[-1].split(".")[0]
@@ -69,7 +69,7 @@ def convex_decomposition_file(
         coacd.set_log_level("error")
 
     # executing the convex decomposition
-    meshfile = _check_filepath(meshfile)
+    meshfile = check_filepath(meshfile)
     _mesh = trimesh.load(meshfile, force="mesh")
     full_mesh = coacd.Mesh(_mesh.vertices, _mesh.faces)
     parts = coacd.run_coacd(full_mesh, **kwargs)  # list of (vert, face) tuples
@@ -102,7 +102,7 @@ def convex_decomposition_dir(
     Returns:
         all_decomposed_meshes: A list of lists of Trimesh objects representing the convex decompositions.
     """
-    meshdir = _check_filepath(meshdir)
+    meshdir = check_filepath(meshdir)
     if recursive:
         glob_func = Path(meshdir).rglob
     else:

ambersim/utils/io_utils.py

@@ -6,27 +6,33 @@
 import numpy as np
 import trimesh
 from dm_control import mjcf
+from dm_control.mjcf.element import _ElementImpl
+from dm_control.mjcf.parser import from_file
+from dm_control.mjcf.physics import Physics
 from mujoco import mjx
+from mujoco.mjx._src.device import device_put
+from mujoco.mjx._src.io import make_data
 
 from ambersim import ROOT
-from ambersim.utils._internal_utils import _check_filepath
+from ambersim.utils._internal_utils import check_filepath
 from ambersim.utils.conversion_utils import save_model_xml
 
 
-def _modify_robot_float_base(filepath: Union[str, Path]) -> mj.MjModel:
+def _modify_robot_float_base(filepath: str) -> mj.MjModel:
     """Modifies a robot to have a floating base if it doesn't already."""
     # loading current robot
     assert str(filepath).split(".")[-1] == "xml"
-    robot = mjcf.from_file(filepath, model_dir="/".join(filepath.split("/")[:-1]))
+    robot = from_file(filepath, model_dir="/".join(filepath.split("/")[:-1]))
 
     # only add free joint if the first body after worldbody has no freejoints
-    assert robot.worldbody is not None
+    assert robot.worldbody is not None  # pyright typechecking
+    assert isinstance(robot.worldbody, _ElementImpl)  # pyright typechecking
     if len(robot.worldbody.body[0].joint) == 0:
         robot.worldbody.body[0].add("freejoint", name="freejoint")
         assert robot.worldbody.body[0].inertial is not None  # checking for non-physical parsing errors
 
     # extracts the mujoco model from the dm_mujoco Physics object
-    physics = mjcf.Physics.from_mjcf_model(robot)
+    physics = Physics.from_mjcf_model(robot)
     assert physics is not None  # pyright typechecking
     model = physics.model.ptr
     return model
@@ -60,7 +66,7 @@ def load_mj_model_from_file(
     elif isinstance(solver, mj.mjtSolver):
         assert solver in [mj.mjtSolver.mjSOL_CG]
 
-    filepath = _check_filepath(filepath)
+    filepath = check_filepath(filepath)
 
     # loading the model and data. check whether freejoint is added forcibly
     if force_float:
@@ -69,7 +75,7 @@ def load_mj_model_from_file(
             output_name = "/".join(str(filepath).split("/")[:-1]) + "/_temp_xml_model.xml"
             save_model_xml(filepath, output_name=output_name)
             mj_model = _modify_robot_float_base(output_name)
-            Path.unlink(output_name)
+            Path(output_name).unlink()
         else:
             mj_model = _modify_robot_float_base(filepath)
     else:
@@ -88,8 +94,8 @@ def load_mj_model_from_file(
 def mj_to_mjx_model_and_data(mj_model: mj.MjModel) -> Tuple[mjx.Model, mjx.Data]:
     """Converts a mujoco model to an mjx (model, data) pair."""
     try:
-        mjx_model = mjx.device_put(mj_model)
-        mjx_data = mjx.make_data(mjx_model)
+        mjx_model = device_put(mj_model)
+        mjx_data = make_data(mjx_model)
         return mjx_model, mjx_data
     except NotImplementedError as e:
         extended_msg = """

examples/convex_decomposition.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 
 from ambersim import ROOT
-from ambersim.utils._internal_utils import _rmtree
+from ambersim.utils._internal_utils import rmtree
 from ambersim.utils.conversion_utils import convex_decomposition_dir, convex_decomposition_file
 
 """This example demonstrates how to perform convex decompositions of nonconvex meshes.
@@ -24,7 +24,7 @@
 print("Example 1: paths of decomposed files:")
 for f in test_save_dir.glob("*.obj"):
     print(str(f))
-_rmtree(test_save_dir)  # remove the test directory (delete this if you want to keep saved files)
+rmtree(test_save_dir)  # remove the test directory (delete this if you want to keep saved files)
 
 # example 2: whole directory
 # setting up dummy model directory with two meshes in it
@@ -44,5 +44,5 @@
 print("Example 2: paths of decomposed files:")
 for f in test_save_dir.glob("*.obj"):
     print(str(f))
-_rmtree(test_model_dir)
-_rmtree(test_save_dir)
+rmtree(test_model_dir)
+rmtree(test_save_dir)

pyrightconfig.json

@@ -1,8 +1,8 @@
 {
   "include": [
-    "ambersim"
+    "ambersim",
   ],
-  "reportMissingImports": false,
+  "reportMissingImports": true,
   "reportMissingTypeStubs": false,
-  "reportGeneralTypeIssues": false,
+  "reportGeneralTypeIssues": true,
 }

tests/test_model_io.py

@@ -8,7 +8,7 @@
 from mujoco import mjx
 
 from ambersim import ROOT
-from ambersim.utils._internal_utils import _rmtree
+from ambersim.utils._internal_utils import rmtree
 from ambersim.utils.conversion_utils import convex_decomposition_file, save_model_xml
 from ambersim.utils.introspection_utils import get_joint_names
 from ambersim.utils.io_utils import _modify_robot_float_base, load_mjx_model_and_data_from_file
@@ -37,7 +37,7 @@ def test_load_model():
     assert load_mjx_model_and_data_from_file(Path(repo_path))
 
     # remove temp local dir
-    _rmtree(local_dir)
+    rmtree(local_dir)
 
 
 def test_all_models():
@@ -126,4 +126,4 @@ def test_convex_decomposition():
         signed_dist_t = igl.signed_distance(coords, tverts, tfaces)[0]
         assert np.allclose(signed_dist_d, signed_dist_t)
 
-    _rmtree(savedir)
+    rmtree(savedir)

typings/coacd/__init__.pyi

@@ -0,0 +1,47 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import ctypes
+import os
+from ctypes import POINTER, Structure, c_bool, c_char_p, c_double, c_int, c_uint, c_uint64, c_void_p
+
+import numpy as np
+import trimesh
+
+_lib_files = ...
+
+class CoACD_Mesh(ctypes.Structure):
+    _fields_ = ...
+
+class CoACD_MeshArray(ctypes.Structure):
+    _fields_ = ...
+
+class Mesh:
+    def __init__(self, vertices=..., indices=...) -> None: ...
+
+def run_coacd(
+    mesh: Mesh,
+    threshold: float = ...,
+    max_convex_hull: int = ...,
+    preprocess_mode: str = ...,
+    preprocess_resolution: int = ...,
+    resolution: int = ...,
+    mcts_nodes: int = ...,
+    mcts_iterations: int = ...,
+    mcts_max_depth: int = ...,
+    pca: int = ...,
+    merge: bool = ...,
+    seed: int = ...,
+):  # -> list[Unknown]:
+    ...
+
+def set_log_level(level: str):  # -> None:
+    ...
+
+if __name__ == "__main__":
+    mesh = ...
+    mesh = ...
+    result = ...
+    mesh_parts = ...
+    scene = ...

typings/dm_control/__init__.pyi

@@ -0,0 +1,7 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+"""
+This type stub file was generated by pyright.
+"""

typings/dm_control/_render/__init__.pyi

@@ -0,0 +1,26 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import collections
+import os
+
+from absl import logging
+from dm_control._render import constants
+
+"""OpenGL context management for rendering MuJoCo scenes.
+
+By default, the `Renderer` class will try to load one of the following rendering
+APIs, in descending order of priority: GLFW > EGL > OSMesa.
+
+It is also possible to select a specific backend by setting the `MUJOCO_GL=`
+environment variable to 'glfw', 'egl', or 'osmesa'.
+"""
+BACKEND = ...
+_ALL_RENDERERS = ...
+_NO_RENDERER = ...
+if BACKEND is not None:
+    import_func = ...
+    Renderer = ...
+else: ...
+USING_GPU = ...

typings/dm_control/_render/base.pyi

@@ -0,0 +1,59 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import abc
+import contextlib
+
+"""Base class for OpenGL context handlers.
+
+`ContextBase` defines a common API that OpenGL rendering contexts should conform
+to. In addition, it provides a `make_current` context manager that:
+
+1. Makes this OpenGL context current within the appropriate rendering thread.
+2. Yields an object exposing a `call` method that can be used to execute OpenGL
+   calls within the rendering thread.
+
+See the docstring for `dm_control.utils.render_executor` for further details
+regarding rendering threads.
+"""
+_CURRENT_CONTEXT_FOR_THREAD = ...
+_CURRENT_THREAD_FOR_CONTEXT = ...
+
+class ContextBase(metaclass=abc.ABCMeta):
+    """Base class for managing OpenGL contexts."""
+
+    def __init__(self, max_width, max_height, render_executor_class=...) -> None:
+        """Initializes this context."""
+        ...
+    def keep_alive(self, obj):  # -> None:
+        ...
+    def dont_keep_alive(self, obj):  # -> None:
+        ...
+    def increment_refcount(self):  # -> None:
+        ...
+    def decrement_refcount(self):  # -> None:
+        ...
+    @property
+    def terminated(self):  # -> bool:
+        ...
+    @property
+    def thread(self):  # -> Thread | None:
+        ...
+    def free(self):  # -> None:
+        """Frees resources associated with this context if its refcount is zero."""
+        ...
+    def __del__(self):  # -> None:
+        ...
+    @contextlib.contextmanager
+    def make_current(self):  # -> Generator[PassthroughRenderExecutor, Any, None]:
+        """Context manager that makes this Renderer's OpenGL context current.
+
+        Yields:
+          An object that exposes a `call` method that can be used to call a
+          function on the dedicated rendering thread.
+
+        Raises:
+          RuntimeError: If this context is already current on another thread.
+        """
+        ...

typings/dm_control/_render/constants.pyi

@@ -0,0 +1,11 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+"""String constants for the rendering module."""
+MUJOCO_GL = ...
+PYOPENGL_PLATFORM = ...
+OSMESA = ...
+GLFW = ...
+EGL = ...
+NO_RENDERER = ...

typings/dm_control/_render/executor/__init__.pyi

@@ -0,0 +1,32 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import os
+
+from dm_control._render.executor.native_mutex.render_executor import NativeMutexOffloadingRenderExecutor
+from dm_control._render.executor.render_executor import (
+    BaseRenderExecutor,
+    OffloadingRenderExecutor,
+    PassthroughRenderExecutor,
+)
+
+"""RenderExecutor executes OpenGL rendering calls on an appropriate thread.
+
+OpenGL calls must be made on the same thread as where an OpenGL context is
+made current on. With GPU rendering, migrating OpenGL contexts between threads
+can become expensive. We provide a thread-safe executor that maintains a
+thread on which an OpenGL context can be kept permanently current, and any other
+threads that wish to render with that context will have their rendering calls
+offloaded to the dedicated thread.
+
+For single-threaded applications, set the `DISABLE_RENDER_THREAD_OFFLOADING`
+environment variable before launching the Python interpreter. This will
+eliminate the overhead of unnecessary thread-switching.
+"""
+_OFFLOAD = ...
+_EXECUTORS = ...
+if _OFFLOAD:
+    RenderExecutor = ...
+else:
+    RenderExecutor = ...