Cleanup WASM bindings generators so they all have a similar interface

PiperOrigin-RevId: 827965451
Change-Id: I657fc4d3a480fbc62124c84fe9ddc0900a91183f

wasm/codegen/generators/enums.py

@@ -14,7 +14,7 @@
 
 """Generates Embind bindings for MuJoCo enums."""
 
-from typing import Mapping
+from typing import Mapping, Optional
 
 from introspect import ast_nodes
 
@@ -38,10 +38,13 @@ class Generator:
     code += ";"
     return code
 
-  def generate(self) -> str:
+  def generate(self) -> list[tuple[str, list[str]]]:
     """Generates all Embind code for the provided enums."""
 
     code = []
     for enum in self.enums.values():
       code.append(self._generate_enum_binding(enum))
-    return "\n\n".join(code) + "\n"
+
+    content = "\n\n".join(code)
+    marker = "// {{ ENUM_BINDINGS }}"
+    return [(marker, [content])]

wasm/codegen/generators/functions.py

@@ -14,7 +14,7 @@
 
 """Generates Embind bindings for MuJoCo functions."""
 
-from typing import List, Mapping
+from typing import List, Mapping, Optional
 
 from introspect import ast_nodes
 
@@ -76,11 +76,14 @@ class Generator:
 
     return result
 
-  def generate(self) -> tuple[str, str]:
+  def generate(self) -> list[tuple[str, list[str]]]:
     """Generates the bindings file for all functions."""
 
     wrapper_functions = self._generate_wrappers()
     function_bindings = self._generate_direct_bindable_functions()
     function_bindings += self._generate_wrapper_bindable_functions()
 
-    return wrapper_functions, function_bindings
+    return [
+        ("// {{ WRAPPER_FUNCTIONS }}", [wrapper_functions]),
+        ("// {{ FUNCTION_BINDINGS }}", [function_bindings]),
+    ]

wasm/codegen/generators/generators_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 from absl.testing import absltest
-
 from introspect import ast_nodes
 
 from wasm.codegen.generators import enums
@@ -51,13 +50,14 @@ class EnumsGeneratorTest(absltest.TestCase):
     .value("ALPHA", ALPHA)
     .value("BETA", BETA);
 
-  enum_<EmptyEnum>("EmptyEnum");
-"""
+  enum_<EmptyEnum>("EmptyEnum");"""
 
-    actual_code = generator.generate()
+    markers_and_content = generator.generate()
+    actual_code = "\n\n".join(markers_and_content[0][1])
 
     self.assertEqual(actual_code, expected_code)
 
+
 class FunctionsGeneratorTest(absltest.TestCase):
 
   def setUp(self):

wasm/codegen/generators/structs.py

@@ -23,17 +23,15 @@ from wasm.codegen.helpers import structs
 class Generator:
   """Generates C++ code for binding and wrapping MuJoCo structs."""
 
-  def __init__(self):
+  def generate(self) -> list[tuple[str, list[str]]]:
+    """Generates C++ header file for binding and wrapping MuJoCo structs."""
+
     # Traverse the introspect dictionary to get the field
     # wrapper/bindings statements set up for each struct
     self.structs_to_bind_data = structs.generate_wasm_bindings(
         constants.STRUCTS_TO_BIND
     )
 
-  def generate_header(
-      self
-  ) -> list[tuple[str, list[Optional[str]]]]:
-    """Generates C++ header file for binding and wrapping MuJoCo structs."""
     autogenned_struct_definitions = []
     markers_and_content = []
 
@@ -60,22 +58,18 @@ class Generator:
         "// {{ AUTOGENNED_STRUCT_DEFINITIONS }}",
         autogenned_struct_definitions,
     ))
-    return markers_and_content
 
-  def generate_source(self) -> list[tuple[str, list[str]]]:
-    """Generates C++ source file for binding and wrapping MuJoCo structs."""
-    constructors = [
-        (
-            f"// INSERT-GENERATED-{struct_data.wrap_name}-CONSTRUCTOR",
-            [struct_data.wrapped_source],
-        )
-        for _, struct_data in self.structs_to_bind_data.items()
-    ]
-    properties = [
-        (
-            f"// INSERT-GENERATED-{struct_data.wrap_name}-BINDINGS",
-            [l.binding for l in struct_data.wrapped_fields],
-        )
-        for _, struct_data in self.structs_to_bind_data.items()
-    ]
-    return constructors + properties
+    for _, struct_data in self.structs_to_bind_data.items():
+      # Bindings
+      markers_and_content.append((
+          f"// INSERT-GENERATED-{struct_data.wrap_name}-BINDINGS",
+          [l.binding for l in struct_data.wrapped_fields],
+      ))
+
+      # Special member functions
+      markers_and_content.append((
+          f"// INSERT-GENERATED-{struct_data.wrap_name}-CONSTRUCTOR",
+          [struct_data.wrapped_source],
+      ))
+
+    return markers_and_content