Format enums with one value per line (#216)

* Format enums with one value per line

Issue #213

2 files changed, 61 insertions, 20 deletions

diff --git a/pycparser/c_generator.py b/pycparser/c_generator.py
index 73e7f1b..a7ee5e6 100644
--- a/pycparser/c_generator.py
+++ b/pycparser/c_generator.py
@@ -135,18 +135,14 @@ class CGenerator(object):
         return ', '.join(visited_subexprs)
 
     def visit_Enum(self, n):
-        s = 'enum'
-        if n.name: s += ' ' + n.name
-        if n.values:
-            s += ' {'
-            for i, enumerator in enumerate(n.values.enumerators):
-                s += enumerator.name
-                if enumerator.value:
-                    s += ' = ' + self.visit(enumerator.value)
-                if i != len(n.values.enumerators) - 1:
-                    s += ', '
-            s += '}'
-        return s
+        return self._generate_struct_union_enum(n, name='enum')
+
+    def visit_Enumerator(self, n):
+        return '{indent}{name} = {value},\n'.format(
+            indent=self._make_indent(),
+            name=n.name,
+            value=self.visit(n.value),
+        )
 
     def visit_FuncDef(self, n):
         decl = self.visit(n.decl)
@@ -268,13 +264,13 @@ class CGenerator(object):
         return '...'
 
     def visit_Struct(self, n):
-        return self._generate_struct_union(n, 'struct')
+        return self._generate_struct_union_enum(n, 'struct')
 
     def visit_Typename(self, n):
         return self._generate_type(n.type)
 
     def visit_Union(self, n):
-        return self._generate_struct_union(n, 'union')
+        return self._generate_struct_union_enum(n, 'union')
 
     def visit_NamedInitializer(self, n):
         s = ''
@@ -289,22 +285,35 @@ class CGenerator(object):
     def visit_FuncDecl(self, n):
         return self._generate_type(n)
 
-    def _generate_struct_union(self, n, name):
-        """ Generates code for structs and unions. name should be either
-            'struct' or union.
+    def _generate_struct_union_enum(self, n, name):
+        """ Generates code for structs, unions, and enums. name should be
+            'struct', 'union', or 'enum'.
         """
+        if name in ('struct', 'union'):
+            members = n.decls
+            body_function = self._generate_struct_union_body
+        else:
+            assert name == 'enum'
+            members = () if n.values is None else n.values.enumerators
+            body_function = self._generate_enum_body
         s = name + ' ' + (n.name or '')
-        if n.decls:
+        if members:
             s += '\n'
             s += self._make_indent()
             self.indent_level += 2
             s += '{\n'
-            for decl in n.decls:
-                s += self._generate_stmt(decl)
+            s += body_function(members)
             self.indent_level -= 2
             s += self._make_indent() + '}'
         return s
 
+    def _generate_struct_union_body(self, members):
+        return ''.join(self._generate_stmt(decl) for decl in members)
+
+    def _generate_enum_body(self, members):
+        # `[:-2] + '\n'` removes the final `,` from the enumerator list
+        return ''.join(self.visit(value) for value in members)[:-2] + '\n'
+
     def _generate_stmt(self, n, add_indent=False):
         """ Generation from a statement node. This method exists as a wrapper
             for individual visit_* methods to handle different treatment of
diff --git a/tests/test_c_generator.py b/tests/test_c_generator.py
index 6d72237..43204c1 100644
--- a/tests/test_c_generator.py
+++ b/tests/test_c_generator.py
@@ -1,4 +1,5 @@
 import sys
+import textwrap
 import unittest
 
 # Run from the root dir
@@ -280,6 +281,37 @@ class TestCtoC(unittest.TestCase):
         self._assert_ctoc_correct('int i = ++(int){ 1 };')
         self._assert_ctoc_correct('struct foo_s foo = (struct foo_s){ 1, 2 };')
 
+    def test_enum(self):
+        s = textwrap.dedent(r'''
+            enum e
+            {
+              a = 1,
+              b = 2,
+              c = 3
+            };
+        '''[1:])
+
+        self._assert_ctoc_correct(s)
+
+        ast = parse_to_ast(s)
+        generator = c_generator.CGenerator()
+        assert generator.visit(ast) == s
+
+    def test_enum_typedef(self):
+        self._assert_ctoc_correct('typedef enum EnumName EnumTypedefName;')
+
+    def test_generate_struct_union_enum_exception(self):
+        generator = c_generator.CGenerator()
+        self.assertRaises(
+            AssertionError,
+            generator._generate_struct_union_enum,
+            n=c_ast.Struct(
+                name='TestStruct',
+                decls=[],
+            ),
+            name='',
+        )
+
 
 if __name__ == "__main__":
     unittest.main()