working enum

src/pwp.py

@@ -2,10 +2,11 @@ import struct
 import typing
 from typing import get_type_hints, get_origin, get_args
 
-
 class PWPError(Exception):
     pass
 
+class uint8:
+    pass
 
 def encode_value(value, typ) -> bytes:
     origin = get_origin(typ)
@@ -61,7 +62,9 @@ def encode_value(value, typ) -> bytes:
     if typ is bool:
         return b"\x01" if value else b"\x00"
 
-    # nested PWP class
+    if typ is uint8:
+        return struct.pack("<B", value)
+
     if hasattr(value, "to_com"):
         return value.to_com()
 
@@ -144,6 +147,9 @@ def decode_value(data: bytes, offset: int, typ):
     if typ is bool:
         return data[offset] != 0, offset + 1
 
+    if typ is uint8:
+        return data[offset], offset + 1
+
     # nested PWP class
     if hasattr(typ, "from_com"):
         return typ.from_com(data[offset:]), len(data)
@@ -163,7 +169,7 @@ def pwp(cls):
         return output
 
     @classmethod
-    def from_com(cls, data):
+    def from_com(data):
         values = {}
         offset = 0
 
@@ -176,4 +182,20 @@ def pwp(cls):
     cls.to_com = to_com
     cls.from_com = from_com
 
+    return cls
+
+def pwp_enum(cls):
+    for index, (name, value) in enumerate(cls.__dict__.items()):
+        if isinstance(value, type) and value.__module__ == cls.__module__:
+            value.__annotations__ = {
+                "_id": uint8,
+                **getattr(value, "__annotations__", {})
+            }
+
+            value._id = index
+
+            value = pwp(value)
+
+            setattr(cls, name, value)
+
     return cls