Add tuple array support to wpistruct dataclasses

subprojects/robotpy-wpiutil/tests/test_struct.py

@@ -179,6 +179,69 @@ def test_user_unpack():
     assert wpistruct.unpack(MyStruct, b"\x02\x00\x00\x00\x01\x00\x00\x60\x40") == v
 
 
+@wpistruct.make_wpistruct(name="VectorStruct")
+@dataclasses.dataclass
+class VectorStruct:
+    data: tuple[wpistruct.double, wpistruct.double, wpistruct.double]
+
+
+def test_user_tuple_array_get_schema():
+    assert wpistruct.getSchema(VectorStruct) == "double data[3]"
+
+
+def test_user_tuple_array_get_size():
+    assert wpistruct.getSize(VectorStruct) == 24
+
+
+def test_user_tuple_array_pack():
+    assert wpistruct.pack(VectorStruct((1.0, 2.0, 3.0))) == (
+        b"\x00\x00\x00\x00\x00\x00\xf0?"
+        b"\x00\x00\x00\x00\x00\x00\x00@"
+        b"\x00\x00\x00\x00\x00\x00\x08@"
+    )
+
+
+def test_user_tuple_array_unpack():
+    assert wpistruct.unpack(
+        VectorStruct,
+        b"\x00\x00\x00\x00\x00\x00\xf0?"
+        b"\x00\x00\x00\x00\x00\x00\x00@"
+        b"\x00\x00\x00\x00\x00\x00\x08@",
+    ) == VectorStruct((1.0, 2.0, 3.0))
+
+
+def test_user_tuple_array_rejects_mixed_types():
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "MixedTuple.value has unsupported tuple type hint: "
+            "tuple fields must be fixed-length and homogeneous"
+        ),
+    ):
+
+        @wpistruct.make_wpistruct
+        @dataclasses.dataclass
+        class MixedTuple:
+            value: tuple[int, float]
+
+
+def test_user_rejects_unsupported_type_with_tuple_in_supported_list():
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "BadField.value is not a wpistruct or does not have a supported type hint "
+            "(supported: bool, int8, uint8, int16, uint16, int, int32, uint32, "
+            "int64, uint64, float, double, or fixed-length homogeneous tuple of "
+            "a supported type)"
+        ),
+    ):
+
+        @wpistruct.make_wpistruct
+        @dataclasses.dataclass
+        class BadField:
+            value: str
+
+
 # def test_user_unpack_into():
 #     v1 = MyStruct(2, True, 3.5)
 #     v2 = MyStruct(3, True, 4.5)

subprojects/robotpy-wpiutil/wpiutil/wpistruct/dataclass.py

@@ -57,9 +57,9 @@ class MyStruct:
             z: wpiutil.struct.double
 
     The types defined in the dataclass can be another WPIStruct compatible class
-    (either builtin or user defined); one of int, bool, or float; or you can
-    use one of the ``wpiutil.wpistruct.[u]int*`` values for explicitly sized
-    integer types.
+    (either builtin or user defined); one of int, bool, or float; a fixed-length
+    homogeneous tuple of those supported types; or you can use one of the
+    ``wpiutil.wpistruct.[u]int*`` values for explicitly sized integer types.
     """
 
     def wrap(cls):
@@ -91,6 +91,33 @@ def wrap(cls):
 }
 
 
+def _get_supported_type_names():
+    supported_names = ", ".join(t.__name__ for t in _type_to_fmt.keys())
+    return f"{supported_names}, or fixed-length homogeneous tuple of a supported type"
+
+
+def _get_fixed_tuple_array_info(cls_name: str, field_name: str, ftype: type):
+    origin = typing.get_origin(ftype)
+    if origin is not tuple:
+        return None
+
+    args = typing.get_args(ftype)
+    if not args or args[-1] is Ellipsis:
+        raise TypeError(
+            f"{cls_name}.{field_name} has unsupported tuple type hint: "
+            "tuple fields must be fixed-length and homogeneous"
+        ) from None
+
+    element_type = args[0]
+    if not all(arg == element_type for arg in args):
+        raise TypeError(
+            f"{cls_name}.{field_name} has unsupported tuple type hint: "
+            "tuple fields must be fixed-length and homogeneous"
+        ) from None
+
+    return element_type, len(args)
+
+
 def _process_class(cls, struct_name: typing.Optional[str]):
     resolved_hints = typing.get_type_hints(cls)
     field_names = [field.name for field in dataclasses.fields(cls)]
@@ -109,6 +136,7 @@ def _process_class(cls, struct_name: typing.Optional[str]):
 
     fmts = []
     schema = []
+    unpackvals = []
     cvvals = []
     vvals = []
     packs = []
@@ -124,9 +152,48 @@ def _process_class(cls, struct_name: typing.Optional[str]):
 
             fmts.append(fmt)
             schema.append(f"{stype} {name}")
+            unpackvals.append(f"arg_{name}")
             cvvals.append(f"arg_{name}")
             vvals.append(f"v.{name}")
 
+        elif array_info := _get_fixed_tuple_array_info(cls_name, name, ftype):
+            element_type, array_len = array_info
+            argn = f"arg_{name}"
+            unpack_args = [f"{argn}_{i}" for i in range(array_len)]
+
+            if element_type in _type_to_fmt:
+                fmt, stype = _type_to_fmt[element_type]
+
+                fmts.append(f"{array_len}{fmt}")
+                schema.append(f"{stype} {name}[{array_len}]")
+                unpackvals.extend(unpack_args)
+                cvvals.append(argn)
+                vvals.append(f"*v.{name}")
+                unpacks.append(f"{argn} = ({', '.join(unpack_args)},)")
+
+            elif hasattr(element_type, "WPIStruct"):
+                typn = f"type_{name}"
+
+                ctx[typn] = element_type
+                ts = wpistruct.getTypeName(element_type)
+                schema.append(f"{ts} {name}[{array_len}]")
+                sz = wpistruct.getSize(element_type)
+                fmts.extend(f"{sz}s" for _ in range(array_len))
+                unpackvals.extend(unpack_args)
+                vvals.append(f"*{argn}")
+                cvvals.append(argn)
+                packs.append(f"{argn} = tuple(wpistruct.pack(i) for i in v.{name})")
+                unpack_exprs = [f"wpistruct.unpack({typn}, {a})" for a in unpack_args]
+                unpacks.append(f"{argn} = ({', '.join(unpack_exprs)},)")
+                # unpackIntos.append(f"wpistruct.unpackInto(v.{name}, {argn})")
+                forEachNested.append(f"wpistruct.forEachNested({typn}, fn)")
+
+            else:
+                raise TypeError(
+                    f"{cls_name}.{name} is not a wpistruct or does not have a supported type hint "
+                    f"(supported: {_get_supported_type_names()})"
+                ) from None
+
         elif hasattr(ftype, "WPIStruct"):
             # nested struct
             argn = f"arg_{name}"
@@ -138,20 +205,21 @@ def _process_class(cls, struct_name: typing.Optional[str]):
             sz = wpistruct.getSize(ftype)
             fmts.append(f"{sz}s")
             vvals.append(argn)
+            unpackvals.append(argn)
             cvvals.append(argn)
             packs.append(f"{argn} = wpistruct.pack(v.{name})")
             unpacks.append(f"{argn} = wpistruct.unpack({typn}, {argn})")
             # unpackIntos.append(f"wpistruct.unpackInto(v.{name}, {argn})")
             forEachNested.append(f"wpistruct.forEachNested({typn}, fn)")
 
         else:
-            supported_names = ", ".join(t.__name__ for t in _type_to_fmt.keys())
             raise TypeError(
                 f"{cls_name}.{name} is not a wpistruct or does not have a supported type hint "
-                f"(supported: {supported_names})"
+                f"(supported: {_get_supported_type_names()})"
             ) from None
 
     s = struct.Struct(f"<{''.join(fmts)}")
+    uvals = ", ".join(unpackvals)
     cvals = ", ".join(cvvals)
     vals = ", ".join(vvals)
 
@@ -195,7 +263,7 @@ def _packInto(v, b):
 
         def _unpack(b):
             try:
-                {cvals} = _s.unpack(b)
+                {uvals} = _s.unpack(b)
                 {unpack_stmts}
                 return cls({cvals})
             except Exception as e: