WIP: Dataclass

binmap/__init__.py

@@ -1,13 +1,21 @@
+import dataclasses
 import struct
-from inspect import Parameter, Signature
+from abc import ABC
+from enum import Enum
+from typing import Dict, NewType, Tuple, Type, TypeVar, Union, get_type_hints
+
+T = TypeVar("T")
 
 
 class BaseDescriptor:
     """Base class for all descriptors
 
     :param name: Variable name"""
 
-    def __init__(self, name):
+    def __set_name__(self, obj, name):
+        self.name = name
+
+    def __init__(self, name=""):
         self.name = name
 
 
@@ -19,7 +27,8 @@ def __get__(self, obj, owner):
         return obj.__dict__[self.name]
 
     def __set__(self, obj, value):
-        struct.pack(obj._datafields[self.name], value)
+        type_hints = get_type_hints(obj)
+        struct.pack(datatypemapping[type_hints[self.name]][1], value)
         obj.__dict__[self.name] = value
 
 
@@ -35,191 +44,171 @@ 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):
     """ConstField descriptor keeps it's value"""
 
     def __set__(self, obj, value):
-        raise AttributeError(f"{self.name} is a constant")
-
-
-class BinmapMetaclass(type):
-    """Metaclass for :class:`Binmap` and all subclasses of :class:`Binmap`.
-
-    :class:`BinmapMetaclass` responsibility is to add for adding variables from
-    _datafields, _enums, _constants and add keyword only parameters.
-
-    _datafields starting with ``_pad`` does't get any instance variable mapped.
-
-    _enums get a variable called _{name} which has the binary data."""
-
-    def __new__(cls, clsname, bases, clsdict):
-        clsobject = super().__new__(cls, clsname, bases, clsdict)
-        keys = clsobject._datafields.keys()
-        sig = Signature(
-            [
-                Parameter(
-                    "binarydata",
-                    Parameter.POSITIONAL_OR_KEYWORD,
-                    default=Parameter.default,
-                )
-            ]
-            + [
-                Parameter(name, Parameter.KEYWORD_ONLY, default=Parameter.default)
-                for name in keys
-            ]
-        )
-        setattr(clsobject, "__signature__", sig)
-        for enum in clsobject._enums:
-            for value, const in clsobject._enums[enum].items():
-                if hasattr(clsobject, const.upper()):
-                    raise ValueError(f"{const} already defined")
-                setattr(clsobject, const.upper(), value)
-        for name in keys:
-            if name.startswith("_pad"):
-                setattr(clsobject, name, PaddingField(name=name))
-            elif name in clsobject._constants:
-                setattr(clsobject, name, ConstField(name=name))
-            elif name in clsobject._enums:
-                setattr(clsobject, name, EnumField(name=name))
-                setattr(clsobject, f"_{name}", BinField(name=f"_{name}"))
-            else:
-                setattr(clsobject, name, BinField(name=name))
-        return clsobject
-
-
-class Binmap(metaclass=BinmapMetaclass):
-    """A class that maps to and from binary data using :py:class:`struct`.
-    All fields in :attr:`binmap.Binmap._datafields` gets a corresponding
-    variable in instance, except variables starting with ``_pad``, which is
-    just padding.
-
-    To create an enum mapping you could add :attr:`binmap.Binmap._enums` with a
-    map to your corresponing datafield:
-
-    .. code-block:: python
-
-        class TempWind(Binmap):
-            _datafields = {"temp": "b", "wind": "B"}
-            _enums = {"wind": {0: "North", 1: "East", 2: "South", 4: "West"}}
-
-        tw = TempWind()
-        tw.temp = 3
-        tw.wind = "South"
-        print(bytes(tw))
-        b'\\x03\\x02'
-
-
-    """
-
-    #: _byteorder: charcter with byteorder
-    _byteorder = ">"
-    #: _datafields: dict with variable name as key and :py:class:`struct` `format strings`_ as value
-    _datafields = {}
-    #: _enums: dict of dicts containing maps of strings
-    _enums = {}
-    #: _constants: dict of constants. This creates a variable that is allways
-    #: the same value. It won't accept binary data with any other value
-    _constants = {}
-
-    def __init__(self, *args, **kwargs):
-        self._formatstring = self._byteorder
-        for fmt in self._datafields.values():
-            self._formatstring += fmt
-
-        bound = self.__signature__.bind(*args, **kwargs)
+        if self.name in obj.__dict__:
+            raise AttributeError(f"{self.name} is a constant")
+        else:
+            obj.__dict__.update({self.name: value})
+
+
+char = NewType("char", int)
+signedchar = NewType("signedchar", int)
+unsignedchar = NewType("unsignedchar", int)
+boolean = NewType("boolean", bool)
+short = NewType("short", int)
+unsignedshort = NewType("unsignedshort", int)
+integer = NewType("integer", int)
+unsignedinteger = NewType("unsignedinteger", int)
+long = NewType("long", int)
+unsignedlong = NewType("unsignedlong", int)
+longlong = NewType("longlong", int)
+unsignedlonglong = NewType("unsignedlonglong", int)
+halffloat = NewType("halffloat", float)
+floating = NewType("floating", float)
+double = NewType("double", float)
+string = NewType("string", str)
+pascalstring = NewType("pascalstring", str)
+pad = NewType("pad", int)
+
+datatypemapping: Dict[type, Tuple[Type[BaseDescriptor], str]] = {
+    char: (BinField, "c"),
+    signedchar: (BinField, "b"),
+    unsignedchar: (BinField, "B"),
+    boolean: (BinField, "?"),
+    bool: (BinField, "?"),
+    short: (BinField, "h"),
+    unsignedshort: (BinField, "H"),
+    integer: (BinField, "i"),
+    int: (BinField, "i"),
+    unsignedinteger: (BinField, "I"),
+    long: (BinField, "l"),
+    unsignedlong: (BinField, "L"),
+    longlong: (BinField, "q"),
+    unsignedlonglong: (BinField, "Q"),
+    halffloat: (BinField, "e"),
+    floating: (BinField, "f"),
+    float: (BinField, "f"),
+    double: (BinField, "d"),
+    string: (BinField, "s"),
+    str: (BinField, "s"),
+    pascalstring: (BinField, "p"),
+    pad: (PaddingField, "x"),
+}
+
+
+def binmapdataclass(cls: Type[T]) -> Type[T]:
+    dataclasses.dataclass(cls)
+    type_hints = get_type_hints(cls)
+
+    cls._formatstring = ""
+
+    for field_ in dataclasses.fields(cls):
+        _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
+        if (
+            type_hints[field_.name] is string
+            or type_hints[field_.name] is pascalstring
+            or type_hints[field_.name] is str
+        ):
+            _type = str(field_.metadata["length"]) + _type
+        cls._formatstring += _type
+
+    return cls
+
+
+def padding(length: int = 1) -> dataclasses.Field:
+    return dataclasses.field(default=length, repr=False, metadata={"padding": True})
+
+
+def constant(value: Union[int, float, str]) -> dataclasses.Field:
+    return dataclasses.field(default=value, init=False, metadata={"constant": True})
+
+
+def stringfield(
+    length: int = 1, default: bytes = b"", fillchar: bytes = b" "
+) -> dataclasses.Field:
+    if default == b"":
+        _default = b"\x00" * length
+    else:
+        _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 = ""
+    _formatstring = ""
+    __binarydata: dataclasses.InitVar[bytes] = b""
+
+    def __init_subclass__(cls, byteorder: str = ">"):
+        cls._byteorder = byteorder
 
-        if "binarydata" in bound.arguments:
-            self._binarydata = bound.arguments["binarydata"]
-            self._unpacker(self._binarydata)
+    def __bytes__(self):
+        return struct.pack(
+            # TODO: use datclass.fields
+            self._byteorder + self._formatstring,
+            *(
+                v
+                for k, v in self.__dict__.items()
+                if k not in ["_byteorder", "_formatstring", "_binarydata"]
+            ),
+        )
 
-        for param in self.__signature__.parameters.values():
-            if param.name == "binarydata":
+    def __post_init__(self, _binarydata: bytes):
+        if _binarydata != b"":
+            self._unpacker(_binarydata)
+        # Kludgy hack to keep order
+        for f in dataclasses.fields(self):
+            if "padding" in f.metadata:
                 continue
-            if param.name in bound.arguments:
-                if param.name in self._constants:
-                    raise AttributeError(f"{param.name} is a constant")
-                setattr(self, param.name, bound.arguments[param.name])
-            elif param.name in self._constants:
-                self.__dict__[param.name] = self._constants[param.name]
-            elif param.name not in self.__dict__:
-                if self._datafields[param.name] in "BbHhIiLlQq":
-                    setattr(self, param.name, 0)
-                elif self._datafields[param.name] in "efd":
-                    setattr(self, param.name, 0.0)
-                elif self._datafields[param.name] == "c":
-                    setattr(self, param.name, b"\x00")
-                else:
-                    setattr(self, param.name, b"")
-
-        if len(args) == 1:
-            self._binarydata = args[0]
-            self._unpacker(self._binarydata)
-        else:
-            self._binarydata = b""
+            if "constant" in f.metadata:
+                self.__dict__.update({f.name: f.default})
+            else:
+                val = getattr(self, f.name)
+                del self.__dict__[f.name]
+                self.__dict__.update({f.name: val})
 
-    def _unpacker(self, value):
-        args = struct.unpack(self._formatstring, value)
+    def _unpacker(self, value: bytes):
+        type_hints = get_type_hints(self)
+        datafieldsmap = {f.name: f for f in dataclasses.fields(self)}
         datafields = [
-            field for field in self._datafields if not field.startswith("_pad")
+            f.name for f in dataclasses.fields(self) if not (type_hints[f.name] is pad)
         ]
+        args = struct.unpack(self._byteorder + self._formatstring, value)
         for arg, name in zip(args, datafields):
-            if name in self._constants:
-                if arg != self._constants[name]:
+            if "constant" in datafieldsmap[name].metadata:
+                if arg != datafieldsmap[name].default:
                     raise ValueError("Constant doesn't match binary data")
-            else:
-                setattr(self, name, arg)
 
-    def __bytes__(self):
-        """packs or unpacks all variables to a binary structure defined by
-        _datafields' format values"""
-        datas = []
-        for var in self._datafields:
-            if not var.startswith("_pad"):
-                if var in self._enums:
-                    datas.append(getattr(self, f"_{var}"))
-                else:
-                    datas.append(getattr(self, var))
-        return struct.pack(self._formatstring, *datas)
-
-    def frombytes(self, value):
+            setattr(self, name, arg)
+
+    def frombytes(self, value: bytes):
         self._unpacker(value)
         self._binarydata = value
-
-    def __eq__(self, other):
-        if self.__signature__ != other.__signature__:
-            return False
-        for field in self._datafields:
-            v1 = getattr(self, field)
-            v2 = getattr(other, field)
-            if v1 != v2:
-                return False
-        return True
-
-    def __str__(self):
-        retval = f"{self.__class__.__name__}"
-        if self._datafields:
-            for key in self._datafields:
-                if not key.startswith("_pad"):
-                    retval += ", %s=%s" % (key, getattr(self, key))
-        return retval