From 747de1e6aaebf023445940da03cfbd1db55290dc Mon Sep 17 00:00:00 2001
From: Andrew Goessling <agoessling@gmail.com>
Date: Mon, 28 Aug 2023 14:53:29 -0700
Subject: [PATCH] Correct size of enums in Python.

---
 src/c_stuff_sack.py        | 1 +
 src/py_stuff_sack.py       | 9 +--------
 test/test_c_stuff_sack.c   | 3 +++
 test/test_py_stuff_sack.py | 6 ++++++
 4 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/c_stuff_sack.py b/src/c_stuff_sack.py
index 4501b9d..d66fabf 100644
--- a/src/c_stuff_sack.py
+++ b/src/c_stuff_sack.py
@@ -367,6 +367,7 @@ def static_assert(all_types, include_enums=True):
     if isinstance(t, (ss.Primitive, ss.Bitfield)):
       s += f'static_assert(sizeof({c_type_name(t)}) == {t.bytes}, "{c_type_name(t)} size mismatch.");\n'
     if include_enums and isinstance(t, ss.Enum):
+      s += f'static_assert(sizeof({t.name}) == sizeof(int), "Size of {t.name} not equal to int.");\n'
       s += f'static_assert(kNum{t.name} < {" * ".join(["256"] * t.bytes)} / 2, "{t.name} size mismatch.");\n'
     if isinstance(t, ss.Struct):
       for f in t.fields:
diff --git a/src/py_stuff_sack.py b/src/py_stuff_sack.py
index bcb7935..9992159 100644
--- a/src/py_stuff_sack.py
+++ b/src/py_stuff_sack.py
@@ -141,13 +141,6 @@ def to_ctypes_name(name, default=None):
 
 
 def get_globals(library, message_spec):
-  ENUM_MAP = {
-      1: ctypes.c_int8,
-      2: ctypes.c_int16,
-      4: ctypes.c_int32,
-      8: ctypes.c_int64,
-  }
-
   BITFIELD_MAP = {
       1: ctypes.c_uint8,
       2: ctypes.c_uint16,
@@ -193,7 +186,7 @@ def get_globals(library, message_spec):
           'description': t.description,
           'value_descriptions':value_descr,
       }
-      e = _EnumType(t.name, (EnumMixin, ENUM_MAP[t.bytes]), attrs)
+      e = _EnumType(t.name, (EnumMixin, ctypes.c_int), attrs)
       global_vars[t.name] = e
       ctypes_map[t.name] = global_vars[t.name]
 
diff --git a/test/test_c_stuff_sack.c b/test/test_c_stuff_sack.c
index c0dfab3..f2c0d98 100644
--- a/test/test_c_stuff_sack.c
+++ b/test/test_c_stuff_sack.c
@@ -48,6 +48,9 @@ static void TestEnum(void) {
   TEST_ASSERT_EQUAL_INT(128, kNumEnum2Bytes);
   TEST_ASSERT_EQUAL_INT(8, SS_ENUM2_BYTES_TEST_PACKED_SIZE);
 
+  TEST_ASSERT_EQUAL_INT(sizeof(int), sizeof(Enum1Bytes));
+  TEST_ASSERT_EQUAL_INT(sizeof(int), sizeof(Enum2Bytes));
+
   Enum2BytesTest enum_test = {
       .enumeration = kNumEnum2Bytes,
   };
diff --git a/test/test_py_stuff_sack.py b/test/test_py_stuff_sack.py
index a789af5..98e47e4 100644
--- a/test/test_py_stuff_sack.py
+++ b/test/test_py_stuff_sack.py
@@ -1,3 +1,4 @@
+import ctypes
 import os
 import unittest
 
@@ -60,6 +61,11 @@ def test_size(self):
     self.assertEqual(128, len(msg_def.Enum2Bytes))
     self.assertEqual(8, msg_def.Enum2BytesTest.packed_size)
 
+    self.assertEqual(ctypes.sizeof(ctypes.c_int),
+                     ctypes.sizeof(msg_def.Enum1BytesTest().enumeration))
+    self.assertEqual(ctypes.sizeof(ctypes.c_int),
+                     ctypes.sizeof(msg_def.Enum2BytesTest().enumeration))
+
   def test_pack_unpack(self):
     msg, buf = self.get_test_message()