support enum with explicit type (C++11)

Author: andreax79

cstruct/__init__.py

@@ -182,6 +182,7 @@ def parse(
     if __cls__ is None:
         __cls__ = MemCStruct
     cls_def = parse_struct_def(__struct__, __cls__=__cls__, process_muliple_definition=True, **kargs)
+    print('!!!!!', cls_def, kargs)
     if cls_def is None:
         return None
     return cls_def['__cls__'].parse(cls_def, **kargs)

cstruct/abstract.py

@@ -30,10 +30,12 @@
 import hashlib
 from io import StringIO
 from enum import IntEnum, EnumMeta, _EnumDict
-from .base import STRUCTS, ENUMS, DEFAULT_ENUM_SIZE
+import struct
+from .base import STRUCTS, ENUMS, DEFAULT_ENUM_SIZE, ENUM_SIZE_TO_C_TYPE
 from .c_parser import parse_struct, parse_struct_def, parse_enum_def, parse_enum, Tokens
 from .field import calculate_padding, FieldType
-from .exceptions import CStructException
+from .native_types import get_native_type
+from .exceptions import CStructException, ParserError
 
 __all__ = ['CStructMeta', 'AbstractCStruct', 'CEnumMeta', 'AbstractCEnum']
 
@@ -315,11 +317,19 @@ def __prepare__(metacls, cls, bases, **kwds):
 
     def __new__(metacls: Type["CEnumMeta"], cls: str, bases: Tuple[Type, ...], classdict: _EnumDict, **kwds: Any) -> "CEnumMeta":
         inst = super().__new__(metacls, cls, bases, classdict, **kwds)
-
         if len(inst) > 0:
-            if "__size__" not in classdict:
+            if classdict.get("__native_format__"):  # data type specified
+                inst.__size__ = struct.calcsize(classdict["__native_format__"])
+            elif "__size__" in classdict:  # size specified
+                try:
+                    inst.__native_format__ = get_native_type(ENUM_SIZE_TO_C_TYPE[inst.__size__]).native_format
+                except KeyError:
+                    raise ParserError(f"Enum has invalid size. Needs to be in {ENUM_SIZE_TO_C_TYPE.keys()}")
+            else:  # default
                 inst.__size__ = DEFAULT_ENUM_SIZE
+                inst.__native_format__ = get_native_type(ENUM_SIZE_TO_C_TYPE[inst.__size__]).native_format
                 print(f"Warning: __size__ not specified for enum {cls}. Will default to {DEFAULT_ENUM_SIZE} bytes")
+
             if not classdict.get("__anonymous__", False):
                 ENUMS[cls] = inst
         return inst
@@ -341,6 +351,7 @@ def parse(
         __enum__: Union[str, Tokens, Dict[str, Any]],
         __name__: Optional[str] = None,
         __size__: Optional[int] = None,
+        __native_format__: Optional[str] = None,
         **kargs: Dict[str, Any],
     ) -> Type["AbstractCEnum"]:
         """
@@ -350,14 +361,16 @@ def parse(
             __enum__: Definition of the enum in C syntax
             __name__: Name of the new Enum. If empty, a name based on the __enum__ hash is generated
             __size__: Number of bytes that the enum should be read as
+            __native_format__: struct module format
 
         Returns:
             cls: A new class mapping the definition
         """
-
         cls_kargs: Dict[str, Any] = dict(kargs)
         if __size__ is not None:
             cls_kargs['__size__'] = __size__
+        if __native_format__ is not None:
+            cls_kargs['__native_format__'] = __native_format__
 
         if isinstance(__enum__, (str, Tokens)):
             cls_kargs.update(parse_enum_def(__enum__, __cls__=cls, **cls_kargs))

cstruct/c_parser.py

@@ -29,6 +29,7 @@
 from .field import calculate_padding, Kind, FieldType
 from .c_expr import c_eval
 from .exceptions import CStructException, ParserError
+from .native_types import get_native_type
 
 if TYPE_CHECKING:
     from .abstract import AbstractCStruct, AbstractCEnum
@@ -58,6 +59,16 @@ def __init__(self, text: str) -> None:
     def pop(self) -> str:
         return self.tokens.pop(0)
 
+    def pop_c_type(self) -> str:
+        c_type = self.pop()
+        if c_type in ['signed', 'unsigned'] and len(self) > 1:
+            # short int, long int, or long long
+            c_type = c_type + " " + self.pop()
+        elif c_type in ['short', 'long'] and len(self) > 1 and self.get() in ['int', 'long']:
+            # short int, long int, or long long
+            c_type = c_type + " " + self.pop()
+        return c_type
+
     def get(self) -> str:
         return self.tokens[0]
 
@@ -203,11 +214,16 @@ def parse_struct_def(
             if result:
                 result['__cls__'].parse(result, **kargs)
             name = tokens.pop()
+            native_format = None
+            if tokens.get() == ':':  # enumeration type declaration
+                tokens.pop()  # pop ":"
+                type_ = get_native_type(tokens.pop_c_type())
+                native_format = type_.native_format
             if tokens.get() == '{':  # named enum
                 tokens.pop()  # pop "{"
-                result = parse_enum(tokens, __name__=name)
+                result = parse_enum(tokens, __name__=name, native_format=native_format)
             elif name == '{':  # unnamed enum
-                result = parse_enum(tokens)
+                result = parse_enum(tokens, native_format=native_format)
             else:
                 raise ParserError(f"{name} definition expected")
 
@@ -244,9 +260,14 @@ def parse_enum_def(__def__: Union[str, Tokens], **kargs: Any) -> Optional[Dict[s
         raise ParserError(f"enum expected - {kind}")
 
     name = tokens.pop()
+    native_format = None
+    if tokens.get() == ':':  # enumeration type declaration
+        tokens.pop()  # pop ":"
+        type_ = get_native_type(tokens.pop_c_type())
+        native_format = type_.native_format
     if tokens.get() == '{':  # named enum
         tokens.pop()  # pop "{"
-        return parse_enum(tokens, __name__=name)
+        return parse_enum(tokens, __name__=name, native_format=native_format)
     elif name == '{':  # unnamed enum
         return parse_enum(tokens)
     else:
@@ -256,6 +277,7 @@ def parse_enum_def(__def__: Union[str, Tokens], **kargs: Any) -> Optional[Dict[s
 def parse_enum(
     __enum__: Union[str, Tokens],
     __name__: Optional[str] = None,
+    native_format: Optional[str] = None,
     **kargs: Any,
 ) -> Optional[Dict[str, Any]]:
     """
@@ -264,6 +286,7 @@ def parse_enum(
     Args:
         __enum__:       definition of the enum in C syntax
         __name__:       enum name
+        native_format:  struct module format
 
     Returns:
         dict: the parsed definition
@@ -283,7 +306,6 @@ def parse_enum(
             break
 
         name = tokens.pop()
-
         next_token = tokens.pop()
         if next_token in {",", "}"}:  # enum-constant without explicit value
             if len(constants) == 0:
@@ -324,6 +346,7 @@ def parse_enum(
         '__is_union__': False,
         '__is_enum__': True,
         '__name__': __name__,
+        '__native_format__': native_format,
         '__cls__': CEnum,
     }
     return result

cstruct/field.py

@@ -26,7 +26,7 @@
 import struct
 from enum import Enum
 from typing import Optional, Any, List, Type, TYPE_CHECKING
-from .base import NATIVE_ORDER, ENUM_SIZE_TO_C_TYPE
+from .base import NATIVE_ORDER
 from .native_types import get_native_type
 from .exceptions import ParserError
 
@@ -193,10 +193,7 @@ def native_format(self) -> str:
             except KeyError:
                 raise ParserError(f"Unknow type `{self.c_type}`")
         elif self.is_enum:
-            try:
-                return get_native_type(ENUM_SIZE_TO_C_TYPE[self.ref.size]).native_format
-            except KeyError:
-                raise ParserError(f"Enum has invalid size. Needs to be in {ENUM_SIZE_TO_C_TYPE.keys()}")
+            return self.ref.__native_format__
         else:
             return 'c'
 

tests/test_cenum.py

@@ -29,6 +29,12 @@ class HtmlFont(CEnum):
     """
 
 
+class EnumWithType(CEnum):
+    __def__ = """
+        enum EnumWithType : int { a, b, c, d};
+    """
+
+
 class StructWithEnum(cstruct.MemCStruct):
     __byte_order__ = cstruct.LITTLE_ENDIAN
     __def__ = """
@@ -94,7 +100,7 @@ def test_sizeof():
     assert cstruct.sizeof("enum HtmlFont") == 2
 
 
-# def test_type():
-#     color = cstruct.parse("enum Color : short { red, green, blue };")
-#     assert color.__size__ == 2
-#     assert cstruct.sizeof("enum Color") == 2
+def test_type():
+    color = cstruct.parse("enum Color : unsigned short { red, green, blue };")
+    assert color.__size__ == 2
+    assert cstruct.sizeof("enum Color") == 2