Add enum support to cstruct.parse

Author: andreax79

cstruct/__init__.py

@@ -28,7 +28,7 @@
 __date__ = '15 August 2013'
 
 import struct
-from typing import Any, Dict, Optional, Type
+from typing import Any, Dict, Optional, Type, Union
 from .base import (
     LITTLE_ENDIAN,
     BIG_ENDIAN,
@@ -39,7 +39,7 @@
     C_TYPE_TO_FORMAT,
     CHAR_ZERO,
 )
-from .abstract import CStructMeta, AbstractCStruct
+from .abstract import CStructMeta, AbstractCStruct, AbstractCEnum
 from .cstruct import CStruct
 from .c_parser import parse_struct_def
 from .mem_cstruct import MemCStruct
@@ -52,6 +52,7 @@
     'CHAR_ZERO',
     'CStruct',
     'MemCStruct',
+    'CEnum',
     'define',
     'undef',
     'getdef',
@@ -134,13 +135,13 @@ def sizeof(type_: str) -> int:
 
 def parse(
     __struct__: str, __cls__: Optional[Type[AbstractCStruct]] = None, __name__: Optional[str] = None, **kargs: Dict[str, Any]
-) -> Optional[Type[AbstractCStruct]]:
+) -> Union[Type[AbstractCStruct], Type[AbstractCEnum], None]:
     """
-    Return a new class mapping a C struct/union definition.
+    Return a new class mapping a C struct/union/enum definition.
     If the string does not contains any definition, return None.
 
     Args:
-        __struct__ (str): definition of the struct (or union) in C syntax
+        __struct__ (str): definition of the struct (or union/enum) in C syntax
         __cls__ (type): super class - CStruct(default) or MemCStruct
         __name__ (str): name of the new class. If empty, a name based on the __struct__ hash is generated
         __byte_order__ (str): byte order, valid values are LITTLE_ENDIAN, BIG_ENDIAN, NATIVE_ORDER
@@ -154,5 +155,7 @@ def parse(
     cls_def = parse_struct_def(__struct__, __cls__=__cls__, **kargs)
     if cls_def is None:
         return None
+    elif cls_def['__is_enum__']:
+        return AbstractCEnum.parse(cls_def, __name__, **kargs)
     else:
         return __cls__.parse(cls_def, __name__, **kargs)

cstruct/c_parser.py

@@ -31,9 +31,9 @@
 from .exceptions import CStructException, ParserError
 
 if TYPE_CHECKING:
-    from .abstract import AbstractCStruct, AbstractCEnum
+    from .abstract import AbstractCStruct
 
-__all__ = ['parse_struct', 'parse_struct_def', 'Tokens']
+__all__ = ['parse_struct', 'parse_struct_def', 'parse_enum_def', 'Tokens']
 
 
 class Tokens(object):
@@ -48,7 +48,7 @@ def __init__(self, text: str) -> None:
                     _, name, value = line.strip().split(maxsplit=2)
                     DEFINES[name] = c_eval(value)
                 except Exception:
-                    raise ParserError("Parsing line {}".format(line))
+                    raise ParserError(f"Parsing line {line}")
             else:
                 lines.append(line)
         text = " ".join(lines)
@@ -135,9 +135,8 @@ def parse_type(tokens: Tokens, __cls__: Type['AbstractCStruct'], byte_order: Opt
             try:
                 ref = STRUCTS[tail]
             except KeyError:
-                raise ParserError("Unknow {} {}".format(c_type, tail))
+                raise ParserError(f"Unknown '{c_type}' '{tail}'")
     elif c_type.startswith('enum'):
-        # from .abstract import AbstractCEnum
         from .cenum import CEnum
 
         c_type, tail = c_type.split(' ', 1)
@@ -175,8 +174,10 @@ def parse_struct_def(
     if not tokens:
         return None
     kind = tokens.pop()
+    if kind == 'enum':
+        return parse_enum_def(__def__, **kargs)
     if kind not in ['struct', 'union']:
-        raise ParserError("struct or union expected - {}".format(kind))
+        raise ParserError("struct, union, or enum expected - {kind}")
     __is_union__ = kind == 'union'
     vtype = tokens.pop()
     if tokens.get() == '{':  # Named nested struct
@@ -185,7 +186,7 @@ def parse_struct_def(
     elif vtype == '{':  # Unnamed nested struct
         return parse_struct(tokens, __cls__=__cls__, __is_union__=__is_union__, __byte_order__=__byte_order__)
     else:
-        raise ParserError("{} definition expected".format(vtype))
+        raise ParserError(f"{vtype} definition expected")
 
 
 def parse_enum_def(__def__: Union[str, Tokens], **kargs: Any) -> Optional[Dict[str, Any]]:
@@ -242,7 +243,7 @@ def parse_enum(__enum__: Union[str, Tokens], **kargs: Any) -> Optional[Dict[str,
                     break
 
             if len(exp_elems) == 0:
-                raise ParserError(f"enum is missing value expression")
+                raise ParserError("enum is missing value expression")
 
             int_expr = " ".join(exp_elems)
             try:
@@ -259,7 +260,12 @@ def parse_enum(__enum__: Union[str, Tokens], **kargs: Any) -> Optional[Dict[str,
         if next_token == "}":
             break
 
-    result = {'__constants__': constants}
+    result = {
+        '__constants__': constants,
+        '__is_struct__': False,
+        '__is_union__': False,
+        '__is_enum__': True,
+    }
     return result
 
 
@@ -325,7 +331,9 @@ def parse_struct(
         '__fields__': list(fields_types.keys()),
         '__fields_types__': fields_types,
         '__size__': size,
+        '__is_struct__': not __is_union__,
         '__is_union__': __is_union__,
+        '__is_enum__': False,
         '__byte_order__': __byte_order__,
         '__alignment__': max_alignment,
     }