Added enum support

Signed-off-by: Cervenka Dusan <[email protected]>

src/cfile/core.py

@@ -168,6 +168,47 @@ def qualifier(self, name) -> bool:
         else:
             raise KeyError(name)
 
+class EnumMember(Element):
+    """
+    Enum element.
+    """
+
+    def __init__(self,
+                 name: str,
+                 value: int | None) -> None:
+        self.name = name
+        self.value = value
+
+
+class Enum(DataType):
+    """
+    A enum definition
+    """
+
+    def __init__(self, name: str, members: list[EnumMember] | None = None) -> None:
+        super().__init__(name)
+
+        if members == None:
+            self.members = []
+        elif isinstance(members, list):
+            self.members = list(members)
+            enumValue = -1
+            for enum in self.members:
+                if enum.value == None:
+                    enumValue += 1
+                    enum.value = enumValue
+                else:
+                    enumValue = enum.value
+        else:
+            raise TypeError('Invalid argument type for "members"')
+
+    def append(self, member: EnumMember) -> None:
+        """
+        Appends new element to the struct definition
+        """
+        if not isinstance(member, EnumMember):
+            raise TypeError(f'Invalid type, expected "EnumMember", got {str(type(member))}')
+        self.members.append(member)
 
 class StructMember(Element):
     """
@@ -422,6 +463,8 @@ def __init__(self,
             self.expression = "true" if expression else "false"
         elif isinstance(expression, (int, float)):
             self.expression = str(expression)
+        elif isinstance(expression, EnumMember):
+            self.expression = expression.name
         elif isinstance(expression, (str, Element)):
             self.expression = expression
         else:

src/cfile/factory.py

@@ -137,6 +137,22 @@ def type(self,
         """
         return core.Type(type_ref, const, pointer, volatile, array)
 
+    def enum_member(self,
+                    name: str,
+                    value: int | None = None):
+        """
+        New EnumMember
+        """
+        return core.EnumMember(name, value)
+
+    def enum(self,
+             name: str,
+             members: list[core.EnumMember] | None = None):
+        """
+        New Enum
+        """
+        return core.Enum(name, members)
+
     def struct_member(self,
                       name: str,
                       data_type: str | core.Type | core.Struct,

src/cfile/writer.py

@@ -17,17 +17,18 @@ class ElementType(Enum):
     DIRECTIVE = 1
     COMMENT = 2
     TYPE_DECLARATION = 3
-    STRUCT_DECLARATION = 4   # Should this be separate from type declaration?
-    VARIABLE_DECLARATION = 5
-    FUNCTION_DECLARATION = 6
-    TYPEDEF = 7
-    TYPE_INSTANCE = 8
-    STRUCT_INSTANCE = 9
-    VARIABLE_USAGE = 10
-    FUNCTION_CALL = 11
-    STATEMENT = 12
-    BLOCK_START = 13
-    BLOCK_END = 14
+    ENUM_DECLARATION = 4   # Should this be separate from type declaration?
+    STRUCT_DECLARATION = 5   # Should this be separate from type declaration?
+    VARIABLE_DECLARATION = 6
+    FUNCTION_DECLARATION = 7
+    TYPEDEF = 8
+    TYPE_INSTANCE = 9
+    STRUCT_INSTANCE = 10
+    VARIABLE_USAGE = 11
+    FUNCTION_CALL = 12
+    STATEMENT = 13
+    BLOCK_START = 14
+    BLOCK_END = 15
 
 
 class Formatter:
@@ -102,6 +103,7 @@ def __init__(self, style: c_style.StyleOptions) -> None:
         self.switcher_all = {
             "Type": self._write_base_type,
             "TypeDef": self._write_typedef_usage,
+            "Enum": self._write_enum_usage,
             "Struct": self._write_struct_usage,
             "Variable": self._write_variable_usage,
             "Function": self._write_function_usage,
@@ -280,6 +282,8 @@ def _write_declaration(self, elem: core.Declaration) -> None:
             self._write_type_declaration(elem.element)
         elif isinstance(elem.element, core.TypeDef):
             self._write_typedef_declaration(elem.element)
+        elif isinstance(elem.element, core.Enum):
+            self._write_enum_declaration(elem.element)
         elif isinstance(elem.element, core.Struct):
             self._write_struct_declaration(elem.element)
         elif isinstance(elem.element, core.Variable):
@@ -314,6 +318,8 @@ def _write_initializer_member(self, value: Any) -> None:
             self._write(str(value))
         elif isinstance(value, str):
             self._write(f'"{value}"')
+        elif isinstance(value, core.EnumMember):
+            self._write(f'{value.name}')
         else:
             raise NotImplementedError(str(type(value)))
 
@@ -383,6 +389,8 @@ def _write_variable_declaration(self, elem: core.Variable) -> None:
             self._write("extern ")
         if isinstance(elem.data_type, core.Type):
             self._write_type_declaration(elem.data_type)
+        elif isinstance(elem.data_type, core.Enum):
+            self._write_enum_usage(elem.data_type)
         elif isinstance(elem.data_type, core.Struct):
             self._write_struct_usage(elem.data_type)
         elif isinstance(elem.data_type, core.Declaration):
@@ -446,6 +454,8 @@ def _write_typedef_declaration(self, elem: core.TypeDef):
             self._write("const ")
         if isinstance(elem.base_type, core.Type):
             self._write_type_declaration(elem.base_type)
+        elif isinstance(elem.base_type, core.Enum):
+            self._write_enum_usage(elem.base_type)
         elif isinstance(elem.base_type, core.Struct):
             self._write_struct_usage(elem.base_type)
         elif isinstance(elem.base_type, core.Declaration):
@@ -507,6 +517,8 @@ def _write_function_declaration(self, elem: core.Function) -> None:
             self._write("static ")
         if isinstance(elem.return_type, core.Type):
             self._write_type_declaration(elem.return_type)
+        elif isinstance(elem.return_type, core.Enum):
+            self._write_enum_usage(elem.return_type)
         elif isinstance(elem.return_type, core.Struct):
             self._write_struct_usage(elem.return_type)
         else:
@@ -599,6 +611,47 @@ def _write_func_call(self, elem: core.FunctionCall) -> None:
             self._write_expression(arg)
         self._write(")")
 
+    def _write_enum_usage(self, elem: core.Enum) -> None:
+        """
+        Writes enum usage
+        """
+        if not elem.name:
+            raise ValueError("enum doesn't have a name. Did you mean to use a declaration?")
+        self._write(f"enum {elem.name}")
+
+    def _write_enum_declaration(self, elem: core.Enum):
+        """
+        Writes enum declaration
+        """
+        self._write(f"enum {elem.name}")
+        if self.style.brace_wrapping.after_struct:
+            self._eol()
+            self._start_line()
+            self._write("{")
+            self._eol()
+        else:
+            self._write(" {")
+            self._eol()
+        if len(elem.members):
+            self._indent()
+        for member in elem.members:
+            self._start_line()
+            self._write_enum_member(member)
+            self._write(",")
+            self._eol()
+        if len(elem.members):
+            self._dedent()
+        self._start_line()
+        self._write("}")
+        self.last_element = ElementType.ENUM_DECLARATION
+
+    def _write_enum_member(self, elem: core.EnumMember) -> None:
+        """
+        Writes enum member
+        """
+        result = elem.name + " = " + str(elem.value)
+        self._write(result)
+
     def _write_struct_usage(self, elem: core.Struct) -> None:
         """
         Writes struct usage