Various minor cleanups to WASM code generation

* Add newline function to code_builder
* Remove type aliases which made jump to definition less ergonomic
* Add types to constants

PiperOrigin-RevId: 827863199
Change-Id: Ib4390d398926466ecce491a3d0ebe308ed5e0cbe

wasm/codegen/generators/functions.py

@@ -15,26 +15,20 @@
 """Generates Embind bindings for MuJoCo functions."""
 
 import pathlib
-from typing import List, Mapping, TypeAlias
+from typing import List, Mapping
 
 from introspect import ast_nodes
 
 from wasm.codegen.helpers import code_builder
 from wasm.codegen.helpers import function_utils
 
-FunctionDecl: TypeAlias = ast_nodes.FunctionDecl
-FunctionParameterDecl: TypeAlias = ast_nodes.FunctionParameterDecl
-PointerType: TypeAlias = ast_nodes.PointerType
-ValueType: TypeAlias = ast_nodes.ValueType
-Path: TypeAlias = pathlib.Path
-
 
 class Generator:
   """Generates Embind bindings for MuJoCo functions."""
 
-  def __init__(self, functions: Mapping[str, FunctionDecl]):
+  def __init__(self, functions: Mapping[str, ast_nodes.FunctionDecl]):
-    self.direct_bind_functions: List[FunctionDecl] = []
+    self.direct_bind_functions: List[ast_nodes.FunctionDecl] = []
-    self.wrapper_bind_functions: List[FunctionDecl] = []
+    self.wrapper_bind_functions: List[ast_nodes.FunctionDecl] = []
 
     for func in functions.values():
       if function_utils.should_be_wrapped(func):
@@ -63,7 +57,7 @@ class Generator:
     return result
 
   def _generate_function_binding(
-      self, func: FunctionDecl, is_wrapper=False
+      self, func: ast_nodes.FunctionDecl, is_wrapper=False
   ) -> str:
     """Generates the Embind code for a single function."""
 

wasm/codegen/helpers/code_builder.py

@@ -36,6 +36,10 @@ class CodeBuilder:
     else:
       self._lines.append("")
 
+  def newline(self) -> None:
+    """Adds a newline."""
+    self.line("")
+
   def to_string(self) -> str:
     """Returns the complete code string."""
     return "\n".join(self._lines)

wasm/codegen/helpers/common.py

@@ -19,13 +19,11 @@ import pathlib
 
 from wasm.codegen.helpers import constants
 
-Path = pathlib.Path
-
 
 def get_default_output_dir() -> str:
   """Gets the default output directory (sibling of 'generated' folder)."""
   # Get the directory of the current file (generator/base.py)
-  current_dir = Path(__file__).parent
+  current_dir = pathlib.Path(__file__).parent
   # Go up one level to the project root and then down to 'generated'
   default_output_dir = str(current_dir.parent / "generated")
   return default_output_dir

wasm/codegen/helpers/constants.py

@@ -14,9 +14,10 @@
 
 """Constants used in the code generation process."""
 
+from typing import List, Set, Dict
 from introspect import structs as introspect_structs
 
-PRIMITIVE_TYPES = {
+PRIMITIVE_TYPES: Set[str] = {
     # go/keep-sorted start
     "char",
     "double",
@@ -36,7 +37,7 @@ PRIMITIVE_TYPES = {
     # go/keep-sorted end
 }
 
-_PLUGIN_FUNCTIONS = [
+_PLUGIN_FUNCTIONS: List[str] = [
     # go/keep-sorted start
     "mj_getPluginConfig",
     "mj_loadAllPluginLibraries",
@@ -61,7 +62,7 @@ _PLUGIN_FUNCTIONS = [
 ]
 
 # Functions that are bound as class methods
-_CLASS_METHODS = [
+_CLASS_METHODS: List[str] = [
     # go/keep-sorted start
     "mj_compile",
     "mj_copyData",
@@ -83,12 +84,12 @@ _CLASS_METHODS = [
 ]
 
 # Omitted because not very useful
-_WRITABLE_ERROR = [
+_WRITABLE_ERROR: List[str] = [
     "mj_printSchema",
 ]
 
 # Omitted thread management functions
-_THREAD_FUNCTIONS = [
+_THREAD_FUNCTIONS: List[str] = [
     # go/keep-sorted start
     "mju_bindThreadPool",
     "mju_defaultTask",
@@ -100,7 +101,7 @@ _THREAD_FUNCTIONS = [
 ]
 
 # Omitted asset cache functions
-_ASSET_CACHE_FUNCTIONS = [
+_ASSET_CACHE_FUNCTIONS: List[str] = [
     # go/keep-sorted start
     "mj_clearCache",
     "mj_getCache",
@@ -111,7 +112,7 @@ _ASSET_CACHE_FUNCTIONS = [
 ]
 
 # Omitted Virtual Filesystem (VFS) functions
-_VFS_FUNCTIONS = [
+_VFS_FUNCTIONS: List[str] = [
     # go/keep-sorted start
     "mj_addBufferVFS",
     "mj_addFileVFS",
@@ -122,7 +123,7 @@ _VFS_FUNCTIONS = [
 ]
 
 # Omitted irrelevant visual functions
-_VISUAL_FUNCTIONS = [
+_VISUAL_FUNCTIONS: List[str] = [
     # go/keep-sorted start
     "mjv_averageCamera",
     "mjv_copyData",
@@ -133,7 +134,7 @@ _VISUAL_FUNCTIONS = [
     # go/keep-sorted end
 ]
 
-_MEMORY_FUNCTIONS = [
+_MEMORY_FUNCTIONS: List[str] = [
     # go/keep-sorted start
     "mj_freeLastXML",
     "mj_freeStack",
@@ -159,7 +160,7 @@ _MEMORY_FUNCTIONS = [
     # go/keep-sorted end
 ]
 
-_GETTERS_AND_SETTERS = [
+_GETTERS_AND_SETTERS: List[str] = [
     # go/keep-sorted start
     "mjs_appendFloatVec",
     "mjs_appendIntVec",
@@ -180,14 +181,14 @@ _GETTERS_AND_SETTERS = [
     # go/keep-sorted end
 ]
 
-_UTILITY_FUNCTIONS = [
+_UTILITY_FUNCTIONS: List[str] = [
     # go/keep-sorted start
     "mju_getXMLDependencies",
     # go/keep-sorted end
 ]
 
 # List of functions that should be skipped during the code generation process.
-SKIPPED_FUNCTIONS = (
+SKIPPED_FUNCTIONS: List[str] = (
     _CLASS_METHODS +
     _THREAD_FUNCTIONS +
     _MEMORY_FUNCTIONS +
@@ -203,7 +204,7 @@ SKIPPED_FUNCTIONS = (
 # Functions that require special wrappers to infer sizes and make additional
 # validation checks. These functions are not bound automatically but are
 # written by hand instead.
-BOUNDCHECK_FUNCS = [
+BOUNDCHECK_FUNCS: List[str] = [
     # go/keep-sorted start
     "mj_addM",
     "mj_angmomMat",
@@ -287,7 +288,7 @@ BOUNDCHECK_FUNCS = [
 ]
 
 # List of structs that should be skipped during the code generation process.
-SKIPPED_STRUCTS = [
+SKIPPED_STRUCTS: List[str] = [
     # go/keep-sorted start
     "mjCache",
     "mjSDF",
@@ -309,7 +310,7 @@ SKIPPED_STRUCTS = [
 # Anonymous structs are not defined as independent structs in the MuJoCo
 # codebase, but they are part of other structs. This dictionary is used to
 # handle them as if they were independent structs.
-ANONYMOUS_STRUCTS = {
+ANONYMOUS_STRUCTS: Dict[str, Dict[str, str]] = {
     # go/keep-sorted start
     "mjVisualGlobal": {"parent": "mjVisual", "field_name": "global"},
     "mjVisualHeadlight": {"parent": "mjVisual", "field_name": "headlight"},
@@ -322,14 +323,14 @@ ANONYMOUS_STRUCTS = {
 
 # This list is created by subtracting the skipped structs from the list of all
 # structs and adding the anonymous structs.
-STRUCTS_TO_BIND = list(
+STRUCTS_TO_BIND: List[str] = list(
     (set(introspect_structs.STRUCTS.keys()) - set(SKIPPED_STRUCTS)).union(
         ANONYMOUS_STRUCTS.keys()
     )
 )
 
 # List of structs that do not have a default constructor.
-NO_DEFAULT_CONSTRUCTORS = [
+NO_DEFAULT_CONSTRUCTORS: List[str] = [
     # go/keep-sorted start
     "mjContact",
     "mjSolverStat",
@@ -349,7 +350,7 @@ NO_DEFAULT_CONSTRUCTORS = [
 # List of `mjData` fields where the array size should be obtained from other
 # `mjData` members, instead of from `mjModel` members. This is typically the
 # case for fields that are dynamically allocated during the simulation.
-MJDATA_SIZES = [
+MJDATA_SIZES: List[str] = [
     # go/keep-sorted start
     "contact",
     "efc_AR",
@@ -421,7 +422,7 @@ MJDATA_SIZES = [
 
 # Dictionary where keys are the struct names and the values are lists of the
 # fields that are manually specified in the bindings.cc template file.
-MANUALLY_ADDED_FIELDS_FROM_TEMPLATE = {
+MANUALLY_ADDED_FIELDS_FROM_TEMPLATE: Dict[str, List[str]] = {
     # go/keep-sorted start
     "MjData": ["solver", "timer", "warning", "contact"],
     "MjSpec": ["option", "visual", "stat", "element", "compiler"],
@@ -456,7 +457,7 @@ MANUALLY_ADDED_FIELDS_FROM_TEMPLATE = {
 # When generating the code for these fields, a specific cast to `uint8_t*` is
 # required for embind. This dictionary is used to register those fields and
 # their sizes.
-BYTE_FIELDS = {
+BYTE_FIELDS: Dict[str, Dict[str, str]] = {
     "buffer": {"size": "nbuffer"},
     "arena": {"size": "narena"},
 }
@@ -464,12 +465,12 @@ BYTE_FIELDS = {
 # Boolean flag to enable debug prints during the struct wrapper and binding
 # generation process. When set to `True`, it will print additional information
 # about the steps being executed.
-STRUCT_DEBUG_MODE = False
+STRUCT_DEBUG_MODE: bool = False
 
 # These structs require specific function calls for creation and/or deletion,
 # or some of their fields need to be handled manually for now;
 # making their wrapper constructors/destructors non-trivial.
-HARDCODED_WRAPPER_STRUCTS = [
+HARDCODED_WRAPPER_STRUCTS: List[str] = [
     "MjData",
     "MjModel",
     "MjvScene",

wasm/codegen/helpers/function_utils_test.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TypeAlias
-
 from absl.testing import absltest
 from introspect import ast_nodes
 
@@ -21,37 +19,37 @@ from wasm.codegen.helpers import constants
 from wasm.codegen.helpers import function_utils
 
 
-PrimitiveTypes: TypeAlias = constants.PRIMITIVE_TYPES
-ValueType: TypeAlias = ast_nodes.ValueType
-PointerType: TypeAlias = ast_nodes.PointerType
-ArrayType: TypeAlias = ast_nodes.ArrayType
-FunctionParameterDecl: TypeAlias = ast_nodes.FunctionParameterDecl
-FunctionDecl: TypeAlias = ast_nodes.FunctionDecl
-
-
 class FunctionUtilsTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.struct_type = ValueType("MyStruct")
+    self.struct_type = ast_nodes.ValueType("MyStruct")
-    self.ptr_to_int = PointerType(ValueType("int"))
+    self.ptr_to_int = ast_nodes.PointerType(ast_nodes.ValueType("int"))
-    self.func_ret_ptr_int = FunctionDecl(
+    self.func_ret_ptr_int = ast_nodes.FunctionDecl(
-        "func_pi", PointerType(ValueType("int")), [], "doc"
+        "func_pi", ast_nodes.PointerType(ast_nodes.ValueType("int")), [], "doc"
     )
-    self.func_ret_ptr_struct = FunctionDecl(
+    self.func_ret_ptr_struct = ast_nodes.FunctionDecl(
-        "func_ps", PointerType(ValueType("MyStruct")), [], "doc"
+        "func_ps",
+        ast_nodes.PointerType(ast_nodes.ValueType("MyStruct")),
+        [],
+        "doc",
     )
 
   def test_return_is_value_of_type(self):
     self.assertTrue(
         function_utils.return_is_value_of_type(
-            FunctionDecl("func_i", ValueType("int"), [], "doc"), PrimitiveTypes
+            ast_nodes.FunctionDecl(
+                "func_i", ast_nodes.ValueType("int"), [], "doc"
+            ),
+            constants.PRIMITIVE_TYPES,
         )
     )
     self.assertFalse(
         function_utils.return_is_value_of_type(
-            FunctionDecl("func_s", ValueType("MyStruct"), [], "doc"),
+            ast_nodes.FunctionDecl(
-            PrimitiveTypes,
+                "func_s", ast_nodes.ValueType("MyStruct"), [], "doc"
+            ),
+            constants.PRIMITIVE_TYPES,
         )
     )
 
@@ -72,21 +70,25 @@ class FunctionUtilsTest(absltest.TestCase):
     )
 
   def test_param_is_primitive_value(self):
-    param_prim_val = FunctionParameterDecl("prim_v", ValueType("int"))
+    param_prim_val = ast_nodes.FunctionParameterDecl(
-    param_arr = FunctionParameterDecl(
+        "prim_v", ast_nodes.ValueType("int")
-        "arr_v", ArrayType(ValueType("int"), extents=(10,))
+    )
+    param_arr = ast_nodes.FunctionParameterDecl(
+        "arr_v", ast_nodes.ArrayType(ast_nodes.ValueType("int"), extents=(10,))
     )
 
     self.assertTrue(function_utils.param_is_primitive_value(param_prim_val))
     self.assertFalse(function_utils.param_is_primitive_value(param_arr))
 
   def test_param_is_pointer_to_primitive_value(self):
-    param_ptr_to_prim = FunctionParameterDecl("p_prim", self.ptr_to_int)
+    param_ptr_to_prim = ast_nodes.FunctionParameterDecl(
-    param_arr_of_prim = FunctionParameterDecl(
+        "p_prim", self.ptr_to_int
-        "a_prim", ArrayType(ValueType("int"), extents=(10,))
     )
-    param_ptr_to_struct = FunctionParameterDecl(
+    param_arr_of_prim = ast_nodes.FunctionParameterDecl(
-        name="p_struct", type=PointerType(inner_type=self.struct_type)
+        "a_prim", ast_nodes.ArrayType(ast_nodes.ValueType("int"), extents=(10,))
+    )
+    param_ptr_to_struct = ast_nodes.FunctionParameterDecl(
+        name="p_struct", type=ast_nodes.PointerType(inner_type=self.struct_type)
     )
     self.assertTrue(
         function_utils.param_is_pointer_to_primitive_value(param_ptr_to_prim)
@@ -99,14 +101,14 @@ class FunctionUtilsTest(absltest.TestCase):
     )
 
   def test_param_is_pointer_to_struct(self):
-    param_arr_of_struct = FunctionParameterDecl(
+    param_arr_of_struct = ast_nodes.FunctionParameterDecl(
-        "a_struct", ArrayType(self.struct_type, extents=(5,))
+        "a_struct", ast_nodes.ArrayType(self.struct_type, extents=(5,))
     )
-    param_ptr_to_struct = FunctionParameterDecl(
+    param_ptr_to_struct = ast_nodes.FunctionParameterDecl(
-        "p_struct", PointerType(self.struct_type)
+        "p_struct", ast_nodes.PointerType(self.struct_type)
     )
-    param_ptr_to_ptr = FunctionParameterDecl(
+    param_ptr_to_ptr = ast_nodes.FunctionParameterDecl(
-        "p_ptr", PointerType(self.ptr_to_int)
+        "p_ptr", ast_nodes.PointerType(self.ptr_to_int)
     )
     self.assertTrue(
         function_utils.param_is_pointer_to_struct(param_arr_of_struct)
@@ -119,43 +121,46 @@ class FunctionUtilsTest(absltest.TestCase):
     )
 
   def test_should_be_wrapped_with_primitive_ptr_return(self):
-    func = FunctionDecl(
+    func = ast_nodes.FunctionDecl(
         name="get_data",
-        return_type=PointerType(ValueType("int")),
+        return_type=ast_nodes.PointerType(ast_nodes.ValueType("int")),
         parameters=tuple(),
         doc="Returns int pointer",
     )
     self.assertTrue(function_utils.should_be_wrapped(func))
 
   def test_generate_function_wrapper_for_simple_func(self):
-    func = FunctionDecl(
+    func = ast_nodes.FunctionDecl(
         name="get_id",
-        return_type=ValueType("int"),
+        return_type=ast_nodes.ValueType("int"),
         parameters=tuple(),
         doc="Returns an integer ID",
     )
     result = function_utils.generate_function_wrapper(func)
-    self.assertEqual(result, """int get_id_wrapper()
+    self.assertEqual(
+        result,
+        """int get_id_wrapper()
 {
   return get_id();
-}""")
+}""",
+    )
 
   def test_generate_function_wrapper_checking_param(self):
     parameters = (
-        FunctionParameterDecl(
+        ast_nodes.FunctionParameterDecl(
             name="mat",
-            type=PointerType(
+            type=ast_nodes.PointerType(
-                inner_type=ValueType(name="mjtNum", is_const=True),
+                inner_type=ast_nodes.ValueType(name="mjtNum", is_const=True),
             ),
         ),
-        FunctionParameterDecl(
+        ast_nodes.FunctionParameterDecl(
             name="nr",
-            type=ValueType(name="int"),
+            type=ast_nodes.ValueType(name="int"),
         ),
     )
-    func = FunctionDecl(
+    func = ast_nodes.FunctionDecl(
         name="get_id",
-        return_type=ValueType("int"),
+        return_type=ast_nodes.ValueType("int"),
         parameters=parameters,
         doc="Returns an integer ID",
     )
@@ -170,25 +175,25 @@ class FunctionUtilsTest(absltest.TestCase):
     )
 
   def test_get_params_string_with_struct_ptr(self):
-    param = FunctionParameterDecl(
+    param = ast_nodes.FunctionParameterDecl(
         name="my_struct",
-        type=PointerType(ValueType("mystruct")),
+        type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")),
     )
     result = function_utils.get_params_string((param,))
     self.assertEqual(result, ["Mystruct& my_struct"])
 
   def test_get_params_string_maybe_with_conversion_struct_ptr(self):
-    param = FunctionParameterDecl(
+    param = ast_nodes.FunctionParameterDecl(
         name="s",
-        type=PointerType(ValueType("customstruct")),
+        type=ast_nodes.PointerType(ast_nodes.ValueType("customstruct")),
     )
     result = function_utils.get_params_string_maybe_with_conversion((param,))
     self.assertEqual(result, ["s.get()"])
 
   def test_get_compatible_return_call(self):
-    func = FunctionDecl(
+    func = ast_nodes.FunctionDecl(
         name="noop",
-        return_type=ValueType("void"),
+        return_type=ast_nodes.ValueType("void"),
         parameters=tuple(),
         doc="does nothing",
     )
@@ -196,9 +201,9 @@ class FunctionUtilsTest(absltest.TestCase):
     self.assertEqual(result, "noop()")
 
   def test_get_compatible_return_type(self):
-    func = FunctionDecl(
+    func = ast_nodes.FunctionDecl(
         name="get_name",
-        return_type=PointerType(ValueType("char")),
+        return_type=ast_nodes.PointerType(ast_nodes.ValueType("char")),
         parameters=tuple(),
         doc="returns name",
     )
@@ -206,9 +211,9 @@ class FunctionUtilsTest(absltest.TestCase):
     self.assertEqual(result.strip(), "std::string")
 
   def test_get_converted_struct_to_class(self):
-    func = FunctionDecl(
+    func = ast_nodes.FunctionDecl(
         name="get_struct",
-        return_type=PointerType(ValueType("mystruct")),
+        return_type=ast_nodes.PointerType(ast_nodes.ValueType("mystruct")),
         parameters=tuple(),
         doc="returns struct",
     )

wasm/codegen/helpers/struct_constructor_code_builder.py

@@ -74,14 +74,15 @@ def _build_struct_header_internal(
 
     builder.line(f"{struct_name}* get() const {{ return ptr_; }}")
     builder.line(f"void set({struct_name}* ptr) {{ ptr_ = ptr; }}")
-    builder.line("")
+
+    builder.newline()
     builder.line("private:")
     builder.line(f"{struct_name}* ptr_;")
     if not is_mjs:
       builder.line("bool owned_ = false;")
 
     if is_mjs and fields_with_init:
-      builder.line("")
+      builder.newline()
       builder.line("public:")
       for field in fields_with_init:
         if field.definition:

wasm/codegen/helpers/struct_constructor_code_builder_test.py

@@ -17,12 +17,6 @@ from introspect import ast_nodes
 from wasm.codegen.helpers import struct_constructor_code_builder
 from wasm.codegen.helpers import struct_field_handler
 
-StructFieldDecl = ast_nodes.StructFieldDecl
-ValueType = ast_nodes.ValueType
-PointerType = ast_nodes.PointerType
-ArrayType = ast_nodes.ArrayType
-StructDecl = ast_nodes.StructDecl
-
 
 class StructConstructorCodeBuilderTest(absltest.TestCase):
 
@@ -58,10 +52,10 @@ MjLROpt::~MjLROpt() {
     )
 
   def test_constructor_code_with_fields_with_init(self):
-    field_with_init = StructFieldDecl(
+    field_with_init = ast_nodes.StructFieldDecl(
         name="element",
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name="mjsElement"),
+            inner_type=ast_nodes.ValueType(name="mjsElement"),
         ),
         doc="",
     )

wasm/codegen/helpers/struct_field_code_builder.py

@@ -17,13 +17,10 @@
 from introspect import ast_nodes
 from wasm.codegen.helpers import code_builder
 
-StructFieldDecl = ast_nodes.StructFieldDecl
-ValueType = ast_nodes.ValueType
 
-
+def build_primitive_type_definition(field: ast_nodes.StructFieldDecl) -> str:
-def build_primitive_type_definition(field: StructFieldDecl) -> str:
   """Builds the C++ code for a primitive type field wrapper."""
-  if not isinstance(field.type, ValueType):
+  if not isinstance(field.type, ast_nodes.ValueType):
     raise ValueError(f"{field.type} must be ValueType.")
   builder = code_builder.CodeBuilder()
   # build getter for primitive type field
@@ -36,7 +33,7 @@ def build_primitive_type_definition(field: StructFieldDecl) -> str:
 
 
 def build_memory_view_definition(
-    field: StructFieldDecl, array_size_str: str, ptr_expr: str
+    field: ast_nodes.StructFieldDecl, array_size_str: str, ptr_expr: str
 ) -> str:
   """Builds the C++ code for a pointer type field wrapper."""
   builder = code_builder.CodeBuilder()
@@ -49,7 +46,7 @@ def build_memory_view_definition(
   return builder.to_string()
 
 
-def build_string_field_definition(field: StructFieldDecl) -> str:
+def build_string_field_definition(field: ast_nodes.StructFieldDecl) -> str:
   """Builds the C++ code for a string type field wrapper."""
   builder = code_builder.CodeBuilder()
   with builder.block(f"mjString {field.name}() const"):
@@ -63,7 +60,7 @@ def build_string_field_definition(field: StructFieldDecl) -> str:
 
 
 def build_mjvec_pointer_definition(
-    field: StructFieldDecl, vector_type: str
+    field: ast_nodes.StructFieldDecl, vector_type: str
 ) -> str:
   """Builds the C++ code for a mjVec type field wrapper."""
   ptr_field_expr = f"*(ptr_->{field.name})"
@@ -79,7 +76,7 @@ def build_mjvec_pointer_definition(
 
 
 def build_simple_property_binding(
-    field: StructFieldDecl,
+    field: ast_nodes.StructFieldDecl,
     struct_wrapper_name: str,
     add_setter: bool = False,
     add_return_value_policy_as_ref: bool = False,

wasm/codegen/helpers/struct_field_code_builder_test.py

@@ -17,18 +17,13 @@ from introspect import ast_nodes
 
 from wasm.codegen.helpers import struct_field_code_builder
 
-StructFieldDecl = ast_nodes.StructFieldDecl
-ValueType = ast_nodes.ValueType
-PointerType = ast_nodes.PointerType
-ArrayType = ast_nodes.ArrayType
-
 
 class StructFieldCodeBuilderTest(absltest.TestCase):
 
   def test_primitive_type_definition(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="ngeom",
-        type=ValueType(name="int"),
+        type=ast_nodes.ValueType(name="int"),
         doc="number of geoms",
     )
     self.assertEqual(
@@ -44,10 +39,10 @@ void set_ngeom(int value) {
     )
 
   def test_memory_view_definition(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="geom_rgba",
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name="float"),
+            inner_type=ast_nodes.ValueType(name="float"),
         ),
         doc="rgba when material is omitted",
         array_extent=("ngeom", 4),
@@ -64,10 +59,10 @@ emscripten::val geom_rgba() const {
     )
 
   def test_string_field_definition(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="string_field",
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name="mjString"),
+            inner_type=ast_nodes.ValueType(name="mjString"),
         ),
         doc="rgba when material is omitted",
     )
@@ -86,10 +81,10 @@ void set_string_field(const mjString& value) {
     )
 
   def test_mjvec_pointer_definition(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="vector_field",
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name="mjDoubleVec"),
+            inner_type=ast_nodes.ValueType(name="mjDoubleVec"),
         ),
         doc="",
     )
@@ -104,10 +99,10 @@ mjDoubleVec &vector_field() const {
     )
 
   def test_mjbyte_vec_pointer_definition(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="vector_field",
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name="mjByteVec"),
+            inner_type=ast_nodes.ValueType(name="mjByteVec"),
         ),
         doc="",
     )
@@ -122,9 +117,9 @@ std::vector<uint8_t> &vector_field() const {
     )
 
   def test_simple_property_binding(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="ngeom",
-        type=ValueType(name="int"),
+        type=ast_nodes.ValueType(name="int"),
         doc="number of geoms",
     )
     self.assertEqual(
@@ -135,9 +130,9 @@ std::vector<uint8_t> &vector_field() const {
     )
 
   def test_simple_property_binding_with_setter(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="ngeom",
-        type=ValueType(name="int"),
+        type=ast_nodes.ValueType(name="int"),
         doc="",
     )
     self.assertEqual(
@@ -149,9 +144,9 @@ std::vector<uint8_t> &vector_field() const {
     )
 
   def test_simple_property_binding_with_return_value_policy_as_ref(self):
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name="ngeom",
-        type=ValueType(name="int"),
+        type=ast_nodes.ValueType(name="int"),
         doc="",
     )
     self.assertEqual(

wasm/codegen/helpers/struct_field_handler.py

@@ -22,13 +22,6 @@ from wasm.codegen.helpers import constants
 from wasm.codegen.helpers import struct_field_code_builder
 from wasm.codegen.helpers import structs_wrappers_data
 
-AnonymousStructDecl = ast_nodes.AnonymousStructDecl
-ArrayType = ast_nodes.ArrayType
-PointerType = ast_nodes.PointerType
-StructFieldDecl = ast_nodes.StructFieldDecl
-ValueType = ast_nodes.ValueType
-WrappedFieldData = structs_wrappers_data.WrappedFieldData
-
 debug_print = common.debug_print
 
 
@@ -37,7 +30,7 @@ class StructFieldHandler:
 
   def __init__(
       self,
-      field: StructFieldDecl,
+      field: ast_nodes.StructFieldDecl,
       struct_wrapper_name: str,
   ):
     self.field = field
@@ -53,27 +46,27 @@ class StructFieldHandler:
         )
     )
 
-  def generate(self) -> WrappedFieldData:
+  def generate(self) -> structs_wrappers_data.WrappedFieldData:
     """Generates the C++ definition and binding code for the struct field."""
     field_type = self.field.type
-    if isinstance(field_type, ValueType) and (
+    if isinstance(field_type, ast_nodes.ValueType) and (
         field_type.name in constants.PRIMITIVE_TYPES
         or field_type.name.startswith("mjt")
     ):
       return self._handle_primitive()
-    elif isinstance(field_type, PointerType):
+    elif isinstance(field_type, ast_nodes.PointerType):
       return self._handle_pointer()
-    elif isinstance(field_type, ArrayType):
+    elif isinstance(field_type, ast_nodes.ArrayType):
       return self._handle_array()
-    elif isinstance(field_type, ValueType) and field_type.name.startswith("mj"):
+    elif isinstance(field_type, ast_nodes.ValueType) and field_type.name.startswith("mj"):
       return self._handle_mj_struct()
-    elif isinstance(field_type, AnonymousStructDecl):
+    elif isinstance(field_type, ast_nodes.AnonymousStructDecl):
       return self._handle_anonymous_struct()
     return self._undefined()
 
-  def _handle_primitive(self) -> WrappedFieldData:
+  def _handle_primitive(self) -> structs_wrappers_data.WrappedFieldData:
     """Handles the generation of C++ definition and binding code for primitive fields."""
-    return WrappedFieldData(
+    return structs_wrappers_data.WrappedFieldData(
         definition=(
             struct_field_code_builder.build_primitive_type_definition(
                 self.field
@@ -88,17 +81,17 @@ class StructFieldHandler:
         is_primitive_or_fixed_size=True,
     )
 
-  def _handle_pointer(self) -> WrappedFieldData:
+  def _handle_pointer(self) -> structs_wrappers_data.WrappedFieldData:
     """Handles the generation of C++ definition and binding code for pointer fields."""
-    if not isinstance(self.field.type, PointerType):
+    if not isinstance(self.field.type, ast_nodes.PointerType):
       raise ValueError(
           f"Expected PointerType, got {type(self.field.type)} for field"
           f" {self.field.name}"
       )
-    field_type: PointerType = self.field.type
+    field_type: ast_nodes.PointerType = self.field.type
     inner_type_name = (
         field_type.inner_type.name
-        if isinstance(field_type.inner_type, ValueType)
+        if isinstance(field_type.inner_type, ast_nodes.ValueType)
         else ""
     )
     ptr_field_expr = f"ptr_->{self.field.name}"
@@ -111,9 +104,7 @@ class StructFieldHandler:
     elif self.field.name in constants.BYTE_FIELDS.keys():
       # for byte fields, we need to cast the pointer to uint8_t*
       # so embind can correctly interpret the memory view
-      ptr_field_expr = (
+      ptr_field_expr = f"static_cast<uint8_t*>({ptr_field_expr})"
-          f"static_cast<uint8_t*>({ptr_field_expr})"
-      )
       # for these byte fields, there is no array_extent, so we add the size of
       # in the config file based in the documentation
       extent = (constants.BYTE_FIELDS[self.field.name]["size"],)
@@ -121,7 +112,7 @@ class StructFieldHandler:
           extent, self.struct_wrapper_name, self.field.name
       )
     elif inner_type_name == "mjString":
-      return WrappedFieldData(
+      return structs_wrappers_data.WrappedFieldData(
           definition=struct_field_code_builder.build_string_field_definition(
               self.field
           ),
@@ -133,7 +124,7 @@ class StructFieldHandler:
           ),
       )
     elif inner_type_name.startswith("mj") and inner_type_name.endswith("Vec"):
-      return WrappedFieldData(
+      return structs_wrappers_data.WrappedFieldData(
           definition=struct_field_code_builder.build_mjvec_pointer_definition(
               self.field, inner_type_name
           ),
@@ -164,11 +155,11 @@ class StructFieldHandler:
           and self.struct_wrapper_name
           not in constants.MANUALLY_ADDED_FIELDS_FROM_TEMPLATE.keys()
       ):
-        ptr_field = cast(PointerType, self.field.type)
+        ptr_field = cast(ast_nodes.PointerType, self.field.type)
         wrapper_field_name = common.uppercase_first_letter(
-            cast(ValueType, ptr_field.inner_type).name
+            cast(ast_nodes.ValueType, ptr_field.inner_type).name
         )
-        return WrappedFieldData(
+        return structs_wrappers_data.WrappedFieldData(
             definition=f"{wrapper_field_name} {self.field.name};",
             binding=struct_field_code_builder.build_simple_property_binding(
                 self.field,
@@ -185,7 +176,7 @@ class StructFieldHandler:
         )
         return self._get_manual_definition(comment_type="complex pointer field")
 
-    return WrappedFieldData(
+    return structs_wrappers_data.WrappedFieldData(
         definition=(
             struct_field_code_builder.build_memory_view_definition(
                 self.field, array_size_str, ptr_field_expr
@@ -194,10 +185,10 @@ class StructFieldHandler:
         binding=self.simple_property_binding,
     )
 
-  def _handle_array(self) -> WrappedFieldData:
+  def _handle_array(self) -> structs_wrappers_data.WrappedFieldData:
     """Handles the generation of C++ definition and binding code for array fields."""
     field_type = self.field.type
-    if not isinstance(field_type, ArrayType):
+    if not isinstance(field_type, ast_nodes.ArrayType):
       raise ValueError(
           f"Expected ArrayType, got {type(field_type)} for field"
           f" {self.field.name}"
@@ -205,7 +196,7 @@ class StructFieldHandler:
     inner_type = field_type.inner_type
     size = math.prod(field_type.extents)
 
-    if isinstance(inner_type, ValueType):
+    if isinstance(inner_type, ast_nodes.ValueType):
       if inner_type.name in constants.PRIMITIVE_TYPES:
         ptr_expr = f"ptr_->{self.field.name}"
         if len(field_type.extents) > 1:
@@ -213,7 +204,7 @@ class StructFieldHandler:
           # to a pointer, so embind can correctly interpret the memory
           # view
           ptr_expr = f"reinterpret_cast<{inner_type.name}*>({ptr_expr})"
-        return WrappedFieldData(
+        return structs_wrappers_data.WrappedFieldData(
             definition=(
                 struct_field_code_builder.build_memory_view_definition(
                     self.field, str(size), ptr_expr
@@ -229,17 +220,17 @@ class StructFieldHandler:
         return self._get_manual_definition(comment_type="array field")
 
     debug_print(f"\tNOT IMPLEMENTED ARRAY field: {self.field.name}")
-    return WrappedFieldData(
+    return structs_wrappers_data.WrappedFieldData(
         definition=(
             f"// TODO: NOT IMPLEMENTED ARRAY wrapper for {self.field.name}"
         ),
         binding=f"// TODO: NOT IMPLEMENTED ARRAY binding for {self.field.name}",
     )
 
-  def _handle_mj_struct(self) -> WrappedFieldData:
+  def _handle_mj_struct(self) -> structs_wrappers_data.WrappedFieldData:
     """Handles the generation of C++ definition and binding code for mj struct fields."""
     if (
-        isinstance(self.field.type, ValueType)
+        isinstance(self.field.type, ast_nodes.ValueType)
         and self.field.name not in self.manually_added_fields
         and self.field.type.name in constants.STRUCTS_TO_BIND
     ):
@@ -249,7 +240,7 @@ class StructFieldHandler:
       if self.struct_wrapper_name not in constants.HARDCODED_WRAPPER_STRUCTS:
         wrapper_field_name = common.uppercase_first_letter(self.field.type.name)
         definition = f"{wrapper_field_name} {self.field.name};"
-      return WrappedFieldData(
+      return structs_wrappers_data.WrappedFieldData(
           definition=definition,
           binding=struct_field_code_builder.build_simple_property_binding(
               self.field,
@@ -263,7 +254,7 @@ class StructFieldHandler:
       )
     return self._get_manual_definition(comment_type="struct field")
 
-  def _handle_anonymous_struct(self) -> WrappedFieldData:
+  def _handle_anonymous_struct(self) -> structs_wrappers_data.WrappedFieldData:
     """Handles the generation of C++ definition and binding code for anonymous struct fields."""
 
     anonymous_struct_name = ""
@@ -277,11 +268,11 @@ class StructFieldHandler:
         break
 
     if (
-        isinstance(self.field.type, AnonymousStructDecl)
+        isinstance(self.field.type, ast_nodes.AnonymousStructDecl)
         and self.field.name not in self.manually_added_fields
         and anonymous_struct_name in constants.STRUCTS_TO_BIND
     ):
-      return WrappedFieldData(
+      return structs_wrappers_data.WrappedFieldData(
           binding=struct_field_code_builder.build_simple_property_binding(
               self.field,
               self.struct_wrapper_name,
@@ -294,24 +285,26 @@ class StructFieldHandler:
       )
     return self._get_manual_definition(comment_type="anonymous struct field")
 
-  def _undefined(self) -> WrappedFieldData:
+  def _undefined(self) -> structs_wrappers_data.WrappedFieldData:
     """This function adds a TODO comment for fields that are not handled by this class yet."""
-    return WrappedFieldData(
+    return structs_wrappers_data.WrappedFieldData(
         definition=f"// TODO: UNDEFINED definition for {self.field.name}",
         binding=f"// TODO: UNDEFINED binding for {self.field.name}",
     )
 
-  def _get_manual_definition(self, comment_type: str = "") -> WrappedFieldData:
+  def _get_manual_definition(
+      self, comment_type: str = ""
+  ) -> structs_wrappers_data.WrappedFieldData:
     """Helper method to generate a comment as a definition for manually added fields."""
     if self.field.name in self.manually_added_fields:
-      return WrappedFieldData(
+      return structs_wrappers_data.WrappedFieldData(
           definition=(
               f"// {comment_type} is defined manually. {self.field.name}"
           ),
           binding=self.simple_property_binding,
       )
 
-    return WrappedFieldData(
+    return structs_wrappers_data.WrappedFieldData(
         definition=(
             f"// TODO: Define {comment_type} manually for {self.field.name}"
         ),

wasm/codegen/helpers/struct_field_handler_test.py

@@ -18,19 +18,13 @@ from introspect import ast_nodes
 from wasm.codegen.helpers import struct_field_handler
 
 
-StructFieldDecl = ast_nodes.StructFieldDecl
-ValueType = ast_nodes.ValueType
-PointerType = ast_nodes.PointerType
-ArrayType = ast_nodes.ArrayType
-
-
 class StructFieldHandlerTest(absltest.TestCase):
 
   def test_scalar_field(self):
     """Test that a scalar type field is handled correctly."""
-    field_scalar = StructFieldDecl(
+    field_scalar = ast_nodes.StructFieldDecl(
         name='ngeom',
-        type=ValueType(name='int'),
+        type=ast_nodes.ValueType(name='int'),
         doc='number of geoms',
     )
 
@@ -56,10 +50,10 @@ void set_ngeom(int value) {
 
   def test_pointer_type_field(self):
     """Test that a pointer type field is handled correctly."""
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name='geom_rgba',
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name='float'),
+            inner_type=ast_nodes.ValueType(name='float'),
         ),
         doc='rgba when material is omitted',
         array_extent=('ngeom', 4),
@@ -84,10 +78,10 @@ emscripten::val geom_rgba() const {
 
   def test_pointer_type_field_for_byte_type(self):
     """Test that a pointer type field for a byte type is handled correctly."""
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name='buffer',
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name='void'),
+            inner_type=ast_nodes.ValueType(name='void'),
         ),
         doc='main buffer; all pointers point in it (nbuffer bytes)',
     )
@@ -110,10 +104,10 @@ emscripten::val buffer() const {
 
   def test_pointer_type_field_for_mj_struct(self):
     """Test that a pointer type field for a mj struct is handled correctly."""
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name='element',
-        type=PointerType(
+        type=ast_nodes.PointerType(
-            inner_type=ValueType(name='mjsElement'),
+            inner_type=ast_nodes.ValueType(name='mjsElement'),
         ),
         doc='',
     )
@@ -133,10 +127,10 @@ emscripten::val buffer() const {
 
   def test_array_type_field(self):
     """Test that an array type field is handled correctly."""
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name='gravity',
-        type=ArrayType(
+        type=ast_nodes.ArrayType(
-            inner_type=ValueType(name='mjtNum'),
+            inner_type=ast_nodes.ValueType(name='mjtNum'),
             extents=(3,),
         ),
         doc='gravitational acceleration',
@@ -160,10 +154,10 @@ emscripten::val gravity() const {
 
   def test_array_field_with_multi_dimensional_array(self):
     """Test that multi-dimensional arrays are handled correctly."""
-    field = StructFieldDecl(
+    field = ast_nodes.StructFieldDecl(
         name='multi_dim_array',
-        type=ArrayType(
+        type=ast_nodes.ArrayType(
-            inner_type=ValueType(name='float'),
+            inner_type=ast_nodes.ValueType(name='float'),
             extents=(3, 4),
         ),
         doc='description',

wasm/codegen/helpers/structs_parser.py

@@ -27,20 +27,14 @@ from wasm.codegen.helpers import struct_field_handler
 from wasm.codegen.helpers import structs_wrappers_data
 
 
-WrappedFieldData = structs_wrappers_data.WrappedFieldData
-WrappedStructData = structs_wrappers_data.WrappedStructData
-StructFieldHandler = struct_field_handler.StructFieldHandler
-
-AnonymousStructDecl = ast_nodes.AnonymousStructDecl
-StructFieldDecl = ast_nodes.StructFieldDecl
 debug_print = common.debug_print
 
 introspect_structs = structs.STRUCTS
 
 
 def generate_wasm_bindings(
-    wrapped_structs: Dict[str, WrappedStructData],
+    wrapped_structs: Dict[str, structs_wrappers_data.WrappedStructData],
-) -> Dict[str, WrappedStructData]:
+) -> Dict[str, structs_wrappers_data.WrappedStructData]:
   """Generates WASM bindings for MuJoCo structs."""
 
   for struct_name, wrap_data in wrapped_structs.items():
@@ -49,7 +43,7 @@ def generate_wasm_bindings(
     elif struct_name in constants.ANONYMOUS_STRUCTS:
       anonymous_struct = _get_anonymous_struct_field(struct_name)
       if not anonymous_struct or not isinstance(
-          anonymous_struct.type, AnonymousStructDecl
+          anonymous_struct.type, ast_nodes.AnonymousStructDecl
       ):
         raise RuntimeError(f"Anonymous struct not found: {struct_name}")
       struct_fields = anonymous_struct.type.fields
@@ -58,9 +52,11 @@ def generate_wasm_bindings(
 
     debug_print(f"Wrapping struct: {struct_name}")
 
-    fields_with_init: List[WrappedFieldData] = []
+    fields_with_init: List[structs_wrappers_data.WrappedFieldData] = []
     for field in struct_fields:
-      field_gen = StructFieldHandler(field, wrap_data.wrap_name).generate()
+      field_gen = struct_field_handler.StructFieldHandler(
+          field, wrap_data.wrap_name
+      ).generate()
       # If the struct has at least one non-primitive or fixed size field
       # we avoid shallow copy to avoid uninitialized memory.
       if not field_gen.is_primitive_or_fixed_size:
@@ -90,7 +86,7 @@ def generate_wasm_bindings(
 
 def _get_anonymous_struct_field(
     anonymous_structs_key: str,
-) -> StructFieldDecl | None:
+) -> ast_nodes.StructFieldDecl | None:
   """Looks up the given key in the anonymous_structs dict and generates bindings for its fields."""
   info = constants.ANONYMOUS_STRUCTS[anonymous_structs_key]
   parent_decl = introspect_structs[info["parent"]]
@@ -101,7 +97,7 @@ def _get_anonymous_struct_field(
           if hasattr(f, "name")
           and f.name == info["field_name"]
           and hasattr(f, "type")
-          and isinstance(f.type, AnonymousStructDecl)
+          and isinstance(f.type, ast_nodes.AnonymousStructDecl)
       ),
       None,
   )