diff --git a/mjx/mujoco/mjx/codegen/generate_warp_types.py b/mjx/mujoco/mjx/codegen/generate_warp_types.py index 4e13054b..98b57215 100644 --- a/mjx/mujoco/mjx/codegen/generate_warp_types.py +++ b/mjx/mujoco/mjx/codegen/generate_warp_types.py @@ -312,15 +312,6 @@ _NESTED_DATACLASS_MANUAL_METHODS = { """, } -_NESTED_DATACLASS_MANUAL_METHOD_NOTES = { - 'TileSet': ( - ' # Manually kept in this generated shim until TileSet method ' - 'generation is\n' - ' # needed more broadly. Keep this in sync with ' - 'mujoco_warp._src.types.TileSet.\n' - ), -} - def write_nested_dataclass(target_fpath: epath.Path, cls: Any): new_class_body = _build_new_class_body_ast( @@ -343,19 +334,6 @@ class {cls.__name__}: ''') -def _write_manual_method_notes(target_fpath: epath.Path): - """Restores method comments stripped by AST-based file rewrites.""" - src = target_fpath.read_text() - for cls_name, note in _NESTED_DATACLASS_MANUAL_METHOD_NOTES.items(): - class_start = src.index(f'class {cls_name}:') - method_start = src.index( - ' def __eq__(self, other) -> bool:\n', class_start - ) - if note not in src[class_start:method_start]: - src = src[:method_start] + note + src[method_start:] - target_fpath.write_text(src) - - def _get_meta_fields(cls_name: str) -> Set[str]: """Returns the set of fields that should be meta-fields in the pytree.""" m = mujoco.MjModel.from_xml_string(_DUMMY_XML) @@ -603,7 +581,6 @@ def main(argv): write_ndim_annotations(target_fpath) write_nworld_leading_dim(target_fpath) - _write_manual_method_notes(target_fpath) file.write_license(target_fpath) file.format_file(target_fpath)