WIP: Dataclass

binmap/__init__.py

@@ -1,6 +1,7 @@
 import dataclasses
 import struct
 from abc import ABC
+from enum import Enum
 from typing import Dict, NewType, Tuple, Type, TypeVar, Union, get_type_hints
 
 T = TypeVar("T")
@@ -28,7 +29,6 @@ def __get__(self, obj, owner):
     def __set__(self, obj, value):
         type_hints = get_type_hints(obj)
         struct.pack(datatypemapping[type_hints[self.name]][1], value)
-        # struct.pack(obj._datafields[self.name], value)
         obj.__dict__[self.name] = value
 
 
@@ -44,25 +44,18 @@ def __set__(self, obj, value):
         pass
 
 
-class EnumField(BaseDescriptor):
+class EnumField(BinField):
     """EnumField descriptor uses "enum" to map to and from strings. Accepts
     both strings and values when setting. Only values that has a corresponding
     string is allowed."""
 
-    def __get__(self, obj, owner):
-        value = obj.__dict__[f"_{self.name}"]
-        return obj._enums[self.name][value]
-
     def __set__(self, obj, value):
-        if value in obj._enums[self.name]:
-            obj.__dict__[f"_{self.name}"] = value
+        datafieldsmap = {f.name: f for f in dataclasses.fields(obj)}
+        if type(value) is str:
+            datafieldsmap[self.name].metadata["enum"][value]
         else:
-            for k, v in obj._enums[self.name].items():
-                if v == value:
-                    obj.__dict__[f"_{self.name}"] = k
-                    return
-
-            raise ValueError("Unknown enum or value")
+            datafieldsmap[self.name].metadata["enum"](value)
+        obj.__dict__[self.name] = value
 
 
 class ConstField(BinField):
@@ -130,6 +123,8 @@ def binmapdataclass(cls: Type[T]) -> Type[T]:
         _base, _type = datatypemapping[type_hints[field_.name]]
         if "constant" in field_.metadata:
             _base = ConstField
+        elif "enum" in field_.metadata:
+            _base = EnumField
         setattr(cls, field_.name, _base(name=field_.name))
         if type_hints[field_.name] is pad:
             _type = field_.default * _type
@@ -153,15 +148,19 @@ def constant(value: Union[int, float, str]) -> dataclasses.Field:
 
 
 def stringfield(
-    length: int = 1, default: str = "", fillchar: str = " "
+    length: int = 1, default: bytes = b"", fillchar: bytes = b" "
 ) -> dataclasses.Field:
-    if default == "":
-        _default = "\x00" * length
+    if default == b"":
+        _default = b"\x00" * length
     else:
-        _default = f"{default:{fillchar}<{length}}"
+        _default = bytes(f"{default:{fillchar}<{length}}")
     return dataclasses.field(default=_default, metadata={"length": length})
 
 
+def enumfield(enumclass: Enum, default: Enum = None) -> dataclasses.Field:
+    return dataclasses.field(default=default, metadata={"enum": enumclass})
+
+
 @dataclasses.dataclass
 class BinmapDataclass(ABC):
     _byteorder = ""

tests/test_binmap.py

@@ -1,4 +1,5 @@
 import struct
+from enum import IntEnum
 
 import pytest
 
@@ -245,56 +246,46 @@ def test_advanced_pack_data(self):
         assert bytes(p) == b"\n\x00\x00<\x00\x00\x00\x00"
 
 
+class WindEnum(IntEnum):
+    North = 0
+    East = 1
+    South = 2
+    West = 3
+
+
 @binmap.binmapdataclass
 class EnumClass(binmap.BinmapDataclass):
     temp: binmap.unsignedchar = 0
-    # wind: int = binmap.enum("h", {0: "North", 1: "East", 2: "South", 4: "West"})
+    wind: binmap.unsignedchar = binmap.enumfield(WindEnum, default=WindEnum.East)
 
 
-@pytest.mark.skip(reason="Enums need to be redesigned")
 class TestEnumClass:
     def test_create_class(self):
         ec = EnumClass()
         assert ec
-        assert EnumClass.SOUTH == 2
 
     def test_get_enum(self):
         ec = EnumClass(temp=10, wind=2)
-        assert ec.wind == "South"
-        assert str(ec) == "EnumClass, temp=10, wind=South"
+        assert ec.wind == WindEnum.South
+        assert str(ec) == "EnumClass(temp=10, wind=2)"
 
     def test_enum_binary(self):
         ec = EnumClass(b"\x0a\x02")
-        assert ec.wind == "South"
-        assert str(ec) == "EnumClass, temp=10, wind=South"
+        assert ec.wind == WindEnum.South
+        assert str(ec) == "EnumClass(temp=10, wind=2)"
 
     def test_set_named_enum(self):
         ec = EnumClass()
-        ec.wind = "South"
+        ec.wind = WindEnum.South
         assert bytes(ec) == b"\x00\x02"
 
-        with pytest.raises(ValueError) as excinfo:
+        with pytest.raises(KeyError) as excinfo:
             ec.wind = "Norhtwest"
-        assert "Unknown enum or value" in str(excinfo)
+        assert "'Norhtwest'" in str(excinfo)
 
         with pytest.raises(ValueError) as excinfo:
             ec.wind = 1.2
-        assert "Unknown enum or value" in str(excinfo)
-
-    def test_colliding_enums(self):
-        with pytest.raises(ValueError) as excinfo:
-
-            class EnumCollide(binmap.BinmapDataclass):
-                _datafields = {
-                    "wind1": "B",
-                    "wind2": "B",
-                }
-                _enums = {
-                    "wind1": {0: "North"},
-                    "wind2": {2: "North"},
-                }
-
-        assert "North already defined" in str(excinfo)
+        assert "1.2 is not a valid WindEnum" in str(excinfo)
 
 
 @binmap.binmapdataclass