Adding type suppressions for pyrefly

PiperOrigin-RevId: 950703735
Change-Id: I6b6d5826a0dd1f4c0a1cbb2a46aee944cd60d884

python/mujoco/__init__.py

@@ -232,8 +232,8 @@ def _bind_data(
 
 _specs.MjSpec.from_zip = from_zip
 _specs.MjSpec.to_zip = to_zip
-_structs.MjData.bind = _bind_data
-_structs.MjModel.bind = _bind_model
+_structs.MjData.bind = _bind_data  # pyrefly: ignore[bad-assignment]
+_structs.MjModel.bind = _bind_model  # pyrefly: ignore[bad-assignment]
 
 HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
 PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')

python/mujoco/codegen/generate_spec_bindings.py

@@ -142,10 +142,10 @@ def _struct_binding_code(
       for f in field.fields
   ):
     for subfield in field.fields:
-      code += _binding_code(subfield, name)
+      code += _binding_code(subfield, name)  # pyrefly: ignore[bad-argument-type]
   # generate for the struct itself
-  field = ast_nodes.ValueType(name=name)
-  code += _value_binding_code(field, classname, varname)
+  field = ast_nodes.ValueType(name=name)  # pyrefly: ignore[bad-assignment]
+  code += _value_binding_code(field, classname, varname)  # pyrefly: ignore[bad-argument-type]
   return code
 
 
@@ -211,7 +211,7 @@ def _ptr_binding_code(
   if vartype == 'mjsElement':  # this is ignored by the caller
     return 'mjsElement'
   if vartype.startswith('mjs'):  # for structs, use the value case
-    return _value_binding_code(field.inner_type, classname, varname)
+    return _value_binding_code(field.inner_type, classname, varname)  # pyrefly: ignore[bad-argument-type]
   elif vartype == 'mjString':  # C++ string -> Python string
     return f"""\
   {classname}.def_property(
@@ -347,7 +347,7 @@ def generate() -> None:
     ) and key != 'mjsElement':
       print('\n  // ' + key)
       for field in structs.STRUCTS[key].fields:
-        code = _binding_code(field, key)
+        code = _binding_code(field, key)  # pyrefly: ignore[bad-argument-type]
         if code != 'mjsElement':
           print(code)
 
@@ -556,7 +556,7 @@ def generate_add() -> None:
       py_args = ['py::arg("name") = py::none()']
 
     for field in structs.STRUCTS[key].fields:
-      line, set_type, name, type_name, cpp_arg, py_arg = _field(field)
+      line, set_type, name, type_name, cpp_arg, py_arg = _field(field)  # pyrefly: ignore[bad-argument-type]
       if line:
         code_field = code_field + '\n        ' + line
         set_types.append(set_type)

python/mujoco/introspect/codegen/generate.py

@@ -126,22 +126,22 @@ def main(argv: Sequence[str]) -> None:
     raise app.UsageError('Too many command-line arguments.')
 
   visitor = ast_processor.process(
-      _JSON_PATH.value,
+      _JSON_PATH.value,  # pyrefly: ignore[bad-argument-type]
       (_HEADER_PATHS.value or '').split(),
       _EXCLUDED,
   )
 
-  with open(_OUT_FUNCTIONS.value, 'w') as f:
+  with open(_OUT_FUNCTIONS.value, 'w') as f:  # pyrefly: ignore[no-matching-overload]
     functions_str = formatter.format_as_python_code(visitor.exported_functions)
     f.write(_HEADER_TEMPLATE.format(year=2022, type='functions'))
     f.write(_FUNCTIONS_TEMPLATE.format(functions_str=functions_str))
 
-  with open(_OUT_ENUMS.value, 'w') as f:
+  with open(_OUT_ENUMS.value, 'w') as f:  # pyrefly: ignore[no-matching-overload]
     enums_str = formatter.format_as_python_code(visitor.exported_enums)
     f.write(_HEADER_TEMPLATE.format(year=2022, type='enums'))
     f.write(_ENUMS_TEMPLATE.format(enums_str=enums_str))
 
-  with open(_OUT_STRUCTS.value, 'w') as f:
+  with open(_OUT_STRUCTS.value, 'w') as f:  # pyrefly: ignore[no-matching-overload]
     structs_str = formatter.format_as_python_code(visitor.exported_structs)
     f.write(_HEADER_TEMPLATE.format(year=2023, type='structs'))
     f.write(_STRUCTS_TEMPLATE.format(structs_str=structs_str))

python/mujoco/introspect/type_parsing.py

@@ -57,7 +57,7 @@ def _parse_maybe_array(
         int(s.strip()) for s in ARRAY_N_PATTERN.findall(array_match.group(0)))
     inner_type_str = type_name[:array_match.start()]
     return ast_nodes.ArrayType(
-        inner_type=_parse_maybe_pointer(inner_type_str.strip(), innermost_type),
+        inner_type=_parse_maybe_pointer(inner_type_str.strip(), innermost_type),  # pyrefly: ignore[bad-argument-type]
         extents=extents)
   else:
     return _parse_maybe_pointer(type_name, innermost_type)
@@ -140,7 +140,7 @@ def parse_type(
   result = None
   while type_str_stack:
     try:
-      result = _parse_maybe_array(type_str_stack.pop(), result)
+      result = _parse_maybe_array(type_str_stack.pop(), result)  # pyrefly: ignore[bad-argument-type]
     except AssertionError as e:
       raise ValueError(f'invalid type name {type_name!r}') from e
 

python/mujoco/minimize.py

@@ -405,7 +405,7 @@ def least_squares(
       print(message, file=output)
 
     # Append log to trace, call iter_callback.
-    log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu)
+    log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu)  # pyrefly: ignore[bad-argument-type]
     if verbose >= Verbosity.FULLITER.value:
       log = dataclasses.replace(
           log, residual=r, jacobian=jac, grad=grad / D, step=step

python/mujoco/rendering/classic/renderer.py

@@ -93,7 +93,7 @@ the clause:
     mujoco.mjr_setBuffer(
         mujoco.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context
     )
-    self._mjr_context.readDepthMap = mujoco.mjtDepthMap.mjDEPTH_ZEROFAR
+    self._mjr_context.readDepthMap = mujoco.mjtDepthMap.mjDEPTH_ZEROFAR  # pyrefly: ignore[bad-assignment]
 
     # Default render flags.
     self._depth_rendering = False
@@ -291,10 +291,10 @@ the clause:
       # Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a
       # camera explicitly defined in the model.
       if camera_id == -1:
-        camera.type = mujoco.mjtCamera.mjCAMERA_FREE
+        camera.type = mujoco.mjtCamera.mjCAMERA_FREE  # pyrefly: ignore[bad-assignment]
         mujoco.mjv_defaultFreeCamera(self._model, camera)
       else:
-        camera.type = mujoco.mjtCamera.mjCAMERA_FIXED
+        camera.type = mujoco.mjtCamera.mjCAMERA_FIXED  # pyrefly: ignore[bad-assignment]
 
     scene_option = scene_option or self._scene_option
     mujoco.mjv_updateScene(

python/mujoco/rollout.py

@@ -180,7 +180,7 @@ class Rollout:
       model = [model]  # Use a length 1 list to simplify code below
 
     if not isinstance(data, list):
-      data = [data]  # Use a length 1 list to simplify code below
+      data = [data]  # Use a length 1 list to simplify code below  # pyrefly: ignore[bad-assignment]
 
     # infer nstep, check for incompatibilities
     nstep = _infer_dimension(
@@ -310,9 +310,9 @@ def rollout(
     ValueError: bad shapes or sizes.
   """  # fmt: skip
   if not isinstance(data, list):
-    data = [data]  # Use a length 1 list to simplify code below
+    data = [data]  # Use a length 1 list to simplify code below  # pyrefly: ignore[bad-assignment]
 
-  nthread = len(data) if len(data) > 1 else 0
+  nthread = len(data) if len(data) > 1 else 0  # pyrefly: ignore[bad-argument-type]
 
   # Use a persistent thread pool if requested
   if persistent_pool:

python/mujoco/sysid/_src/optimize.py

@@ -53,7 +53,7 @@ def _warn_if_ill_conditioned(
           residual=f,
           x=x0.reshape(-1, 1),
           r=r0,
-          eps=eps,
+          eps=eps,  # pyrefly: ignore[bad-argument-type]
           n_res=0,
           bounds=[bounds[0].reshape(-1, 1), bounds[1].reshape(-1, 1)],
       )[0],

python/mujoco/viewer.py

@@ -325,8 +325,8 @@ def _reload(
 
 def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
   """Physics loop for the GUI, to be run in a separate thread."""
-  m: mujoco.MjModel = None
-  d: mujoco.MjData = None
+  m: mujoco.MjModel = None  # pyrefly: ignore[bad-assignment]
+  d: mujoco.MjData = None  # pyrefly: ignore[bad-assignment]
   ctrl_noise = np.array([])
   reload = True