feat: implement the proposal in issue #5

Author: MrIndeciso

libdestruct/common/array/array.py

@@ -7,13 +7,20 @@
 from __future__ import annotations
 
 from abc import abstractmethod
+from types import GenericAlias
 
 from libdestruct.common.obj import obj
 
 
 class array(obj):
     """An array of objects."""
 
+    def __class_getitem__(cls, params: tuple) -> GenericAlias:
+        """Support array[c_int, 3] subscript syntax."""
+        if not isinstance(params, tuple):
+            params = (params,)
+        return GenericAlias(cls, params)
+
     @abstractmethod
     def count(self: array) -> int:
         """Return the size of the array."""

libdestruct/common/array/array_field_inflater.py

@@ -9,6 +9,7 @@
 
 from typing import TYPE_CHECKING
 
+from libdestruct.common.array.array import array
 from libdestruct.common.array.linear_array_field import LinearArrayField
 from libdestruct.common.type_registry import TypeRegistry
 
@@ -18,6 +19,7 @@
     from libdestruct.backing.resolver import Resolver
     from libdestruct.common.obj import obj
 
+
 registry = TypeRegistry()
 
 
@@ -32,4 +34,21 @@ def linear_array_field_inflater(
     return field.inflate
 
 
+def _subscripted_array_handler(
+    item: object,
+    args: tuple,
+    owner: tuple[obj, type[obj]] | None,
+) -> Callable[[Resolver], obj] | None:
+    """Handle subscripted array types like array[c_int, 3]."""
+    if len(args) != 2:
+        return None
+    element_type, count = args
+    if not isinstance(count, int):
+        return None
+    field = LinearArrayField(element_type, count)
+    field.item = registry.inflater_for(element_type)
+    return field.inflate
+
+
 registry.register_instance_handler(LinearArrayField, linear_array_field_inflater)
+registry.register_generic_handler(array, _subscripted_array_handler)

libdestruct/common/enum/enum.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+from types import GenericAlias
 from typing import TYPE_CHECKING
 
 from libdestruct.common.obj import obj
@@ -20,6 +21,12 @@
 class enum(obj):
     """A generic enum."""
 
+    def __class_getitem__(cls, params: tuple) -> GenericAlias:
+        """Support enum[MyEnum] and enum[MyEnum, c_short] subscript syntax."""
+        if not isinstance(params, tuple):
+            params = (params,)
+        return GenericAlias(cls, params)
+
     python_enum: type[Enum]
     """The backing Python enum."""
 

libdestruct/common/enum/enum_field_inflater.py

@@ -8,6 +8,8 @@
 
 from typing import TYPE_CHECKING
 
+from libdestruct.c.c_integer_types import c_int
+from libdestruct.common.enum.enum import enum
 from libdestruct.common.enum.int_enum_field import IntEnumField
 from libdestruct.common.type_registry import TypeRegistry
 
@@ -30,4 +32,19 @@ def generic_enum_field_inflater(
     return field.inflate
 
 
+def _subscripted_enum_handler(
+    item: object,
+    args: tuple,
+    owner: tuple[obj, type[obj]] | None,
+) -> Callable[[Resolver], obj] | None:
+    """Handle subscripted enum types like enum[MyEnum] or enum[MyEnum, c_short]."""
+    if not args:
+        return None
+    python_enum = args[0]
+    backing_type = args[1] if len(args) > 1 else c_int
+    field = IntEnumField(python_enum, size=backing_type.size)
+    return field.inflate
+
+
 registry.register_instance_handler(IntEnumField, generic_enum_field_inflater)
+registry.register_generic_handler(enum, _subscripted_enum_handler)

test/scripts/struct_unit_test.py

@@ -7,7 +7,7 @@
 import unittest
 from enum import IntEnum
 
-from libdestruct import c_int, c_long, c_uint, inflater, struct, ptr, ptr_to_self, array_of, enum, enum_of
+from libdestruct import array, c_int, c_long, c_short, c_uint, inflater, struct, ptr, ptr_to_self, array_of, enum, enum_of
 
 
 class StructMemberCollisionTest(unittest.TestCase):
@@ -264,6 +264,96 @@ class TreeNode(struct):
         self.assertEqual(node.data.value, 42)
 
 
+class SubscriptSyntaxTest(unittest.TestCase):
+    """Test the subscript syntax: enum[T], array[T, N], ptr[T]."""
+
+    def test_enum_subscript(self):
+        """enum[MyEnum] works as a type annotation for struct fields."""
+        class Color(IntEnum):
+            RED = 0
+            GREEN = 1
+            BLUE = 2
+
+        class s_t(struct):
+            color: enum[Color]
+
+        memory = (1).to_bytes(4, "little")
+        s = s_t.from_bytes(memory)
+        self.assertEqual(s.color.value, Color.GREEN)
+
+    def test_enum_subscript_custom_backing(self):
+        """enum[MyEnum, c_short] uses a custom backing type."""
+        class Status(IntEnum):
+            OFF = 0
+            ON = 1
+
+        class s_t(struct):
+            status: enum[Status, c_short]
+
+        memory = (1).to_bytes(2, "little")
+        s = s_t.from_bytes(memory)
+        self.assertEqual(s.status.value, Status.ON)
+        from libdestruct import size_of
+        self.assertEqual(size_of(s_t), 2)
+
+    def test_array_subscript(self):
+        """array[c_int, 3] works as a type annotation for struct fields."""
+        class s_t(struct):
+            data: array[c_int, 3]
+
+        memory = b""
+        for v in [10, 20, 30]:
+            memory += v.to_bytes(4, "little")
+
+        s = s_t.from_bytes(memory)
+        self.assertEqual(s.data[0].value, 10)
+        self.assertEqual(s.data[1].value, 20)
+        self.assertEqual(s.data[2].value, 30)
+
+    def test_array_subscript_size(self):
+        """array[c_int, 3] has correct size."""
+        class s_t(struct):
+            data: array[c_int, 3]
+
+        from libdestruct import size_of
+        self.assertEqual(size_of(s_t), 12)
+
+    def test_ptr_subscript(self):
+        """ptr[T] works as a type annotation (already supported)."""
+        class s_t(struct):
+            val: c_int
+            ref: ptr[c_int]
+
+        memory = b""
+        memory += (42).to_bytes(4, "little")
+        memory += (0).to_bytes(8, "little")
+
+        s = s_t.from_bytes(memory)
+        self.assertEqual(s.val.value, 42)
+
+    def test_mixed_subscript_struct(self):
+        """Struct mixing all subscript syntaxes."""
+        class Direction(IntEnum):
+            UP = 0
+            DOWN = 1
+
+        class s_t(struct):
+            dir: enum[Direction]
+            coords: array[c_int, 2]
+            next: ptr["s_t"]
+
+        memory = b""
+        memory += (1).to_bytes(4, "little")          # dir = DOWN
+        memory += (10).to_bytes(4, "little")          # coords[0]
+        memory += (20).to_bytes(4, "little")          # coords[1]
+        memory += (0).to_bytes(8, "little")           # next = null
+
+        s = s_t.from_bytes(memory)
+        self.assertEqual(s.dir.value, Direction.DOWN)
+        self.assertEqual(s.coords[0].value, 10)
+        self.assertEqual(s.coords[1].value, 20)
+
+
 class StructEqualityTest(unittest.TestCase):
     def test_struct_eq_non_struct_returns_not_implemented(self):
         """struct.__eq__ returns NotImplemented for non-struct values."""