From bdaec5758b7e46eb0b1700795fec6bf07e1b73ed Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Tue, 25 Nov 2025 04:53:04 -0800
Subject: [PATCH] Generate MuJoCo Warp documentation with modified types.py.

PiperOrigin-RevId: 836617034
Change-Id: Ia925b88679744267b34504b146fa712d4195eda4
---
 .readthedocs.yml           |  1 +
 doc/mjwarp/update_types.py | 82 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 doc/mjwarp/update_types.py

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 309c2496..2a77f555 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -21,6 +21,7 @@ build:
       # replace mujoco.mjx.third_party.mujoco_warp import paths with mujoco_warp
       - |
         find mjx/mujoco/mjx/third_party/mujoco_warp -type f -exec sed -i 's/mujoco\.mjx\.third_party\.mujoco_warp/mujoco_warp/g' {} \;
+      - python doc/mjwarp/update_types.py mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
       - UV_PROJECT_ENVIRONMENT=$READTHEDOCS_VIRTUALENV_PATH uv pip install mjx/mujoco/mjx/third_party/mujoco_warp
     install:
       - "true"  # skip
diff --git a/doc/mjwarp/update_types.py b/doc/mjwarp/update_types.py
new file mode 100644
index 00000000..a19a0b61
--- /dev/null
+++ b/doc/mjwarp/update_types.py
@@ -0,0 +1,82 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Update types.py with valid attribute types.
+
+Instances like
+
+@dataclasses.dataclass
+class Option:
+  ...
+  timestep: array("*", float)
+  ...
+
+that do not have valid attribute types are not properly parsed during
+documentation generation.
+
+This script updates such instances with valid types
+
+@dataclasses.dataclass
+class Option:
+  ...
+  timestep: wp.array(dtype=float)
+  ...
+"""
+
+import re
+import sys
+
+
+def replace_array_calls(match):
+  """Replaces array() calls with wp.array, wp.array2d, etc."""
+  args_str = match.group(1)
+  if '...' in args_str:
+    return match.group(0)  # Do not replace call in docstring
+
+  args = [a.strip() for a in args_str.split(',')]
+  n_args = len(args)
+  dtype = args[-1]
+
+  if n_args == 2:
+    return f'wp.array(dtype={dtype})'
+  elif n_args == 3:
+    return f'wp.array2d(dtype={dtype})'
+  elif n_args == 4:
+    return f'wp.array3d(dtype={dtype})'
+  elif n_args == 5:
+    return f'wp.array4d(dtype={dtype})'
+  else:
+    return match.group(0)
+
+
+def process_file(filepath):
+  """Reads types.py, replaces array() calls, and writes back to file."""
+  try:
+    with open(filepath, 'r') as f:
+      content = f.read()
+  except FileNotFoundError:
+    print(f'Error: Could not find {filepath}', file=sys.stderr)
+    return
+
+  new_content = re.sub(r'array\(([^)]+)\)', replace_array_calls, content)
+
+  with open(filepath, 'w') as f:
+    f.write(new_content)
+  print(f'Processed {filepath}')
+
+
+if __name__ == '__main__':
+  if len(sys.argv) != 2:
+    print('Usage: python update_types.py <path_to_types.py>', file=sys.stderr)
+    sys.exit(1)
+  process_file(sys.argv[1])