PWP

src/pwp.py

@@ -0,0 +1,179 @@
+import struct
+import typing
+from typing import get_type_hints, get_origin, get_args
+
+
+class PWPError(Exception):
+    pass
+
+
+def encode_value(value, typ) -> bytes:
+    origin = get_origin(typ)
+
+    # Optional[T]
+    if origin is typing.Union:
+        args = get_args(typ)
+
+        if type(None) in args:
+            if value is None:
+                return b"\x00"
+
+            real_type = next(t for t in args if t is not type(None))
+            return b"\x01" + encode_value(value, real_type)
+
+    # list[T]
+    if origin is list:
+        item_type = get_args(typ)[0]
+
+        out = len(value).to_bytes(4, "little")
+
+        for item in value:
+            out += encode_value(item, item_type)
+
+        return out
+
+    # dict[K, V]
+    if origin is dict:
+        key_type, value_type = get_args(typ)
+
+        out = len(value).to_bytes(4, "little")
+
+        for k, v in value.items():
+            out += encode_value(k, key_type)
+            out += encode_value(v, value_type)
+
+        return out
+
+    # primitives
+    if typ is str:
+        data = value.encode("utf-8")
+        return len(data).to_bytes(4, "little") + data
+
+    if typ is bytes:
+        return len(value).to_bytes(4, "little") + value
+
+    if typ is int:
+        return struct.pack("<q", value)
+
+    if typ is float:
+        return struct.pack("<d", value)
+
+    if typ is bool:
+        return b"\x01" if value else b"\x00"
+
+    # nested PWP class
+    if hasattr(value, "to_com"):
+        return value.to_com()
+
+    raise PWPError(f"Unsupported type: {typ}")
+
+
+def decode_value(data: bytes, offset: int, typ):
+    origin = get_origin(typ)
+
+    # Optional[T]
+    if origin is typing.Union:
+        args = get_args(typ)
+
+        if type(None) in args:
+            present = data[offset]
+            offset += 1
+
+            if present == 0:
+                return None, offset
+
+            real_type = next(t for t in args if t is not type(None))
+            return decode_value(data, offset, real_type)
+
+    # list[T]
+    if origin is list:
+        item_type = get_args(typ)[0]
+
+        count = int.from_bytes(data[offset:offset+4], "little")
+        offset += 4
+
+        result = []
+
+        for _ in range(count):
+            value, offset = decode_value(data, offset, item_type)
+            result.append(value)
+
+        return result, offset
+
+    # dict[K,V]
+    if origin is dict:
+        key_type, value_type = get_args(typ)
+
+        count = int.from_bytes(data[offset:offset+4], "little")
+        offset += 4
+
+        result = {}
+
+        for _ in range(count):
+            key, offset = decode_value(data, offset, key_type)
+            value, offset = decode_value(data, offset, value_type)
+            result[key] = value
+
+        return result, offset
+
+    # primitives
+    if typ is str:
+        length = int.from_bytes(data[offset:offset+4], "little")
+        offset += 4
+
+        return (
+            data[offset:offset+length].decode("utf-8"),
+            offset + length,
+        )
+
+    if typ is bytes:
+        length = int.from_bytes(data[offset:offset+4], "little")
+        offset += 4
+
+        return (
+            data[offset:offset+length],
+            offset + length,
+        )
+
+    if typ is int:
+        return struct.unpack_from("<q", data, offset)[0], offset + 8
+
+    if typ is float:
+        return struct.unpack_from("<d", data, offset)[0], offset + 8
+
+    if typ is bool:
+        return data[offset] != 0, offset + 1
+
+    # nested PWP class
+    if hasattr(typ, "from_com"):
+        return typ.from_com(data[offset:]), len(data)
+
+    raise PWPError(f"Unsupported type: {typ}")
+
+
+def pwp(cls):
+    fields = get_type_hints(cls)
+
+    def to_com(self):
+        output = b""
+
+        for name, typ in fields.items():
+            output += encode_value(getattr(self, name), typ)
+
+        return output
+
+    @classmethod
+    def from_com(cls, data):
+        values = {}
+        offset = 0
+
+        for name, typ in fields.items():
+            value, offset = decode_value(data, offset, typ)
+            values[name] = value
+
+        return cls(**values)
+
+    cls.to_com = to_com
+    cls.from_com = from_com
+
+    return cls