From a98724410cb7f77a3eeede698cfa73496bb889f2 Mon Sep 17 00:00:00 2001
From: biggus-developerus
 <74741815+biggus-developerus@users.noreply.github.com>
Date: Sat, 20 Jul 2024 16:54:12 +0400
Subject: [PATCH] sized packing

---
 packer/_types/__init__.py      |  1 +
 packer/_types/sized_packers.py | 32 ++++++++++++++++++++++++++++++++
 tests/test_sized_packing.py    | 31 +++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)
 create mode 100644 packer/_types/sized_packers.py
 create mode 100644 tests/test_sized_packing.py

diff --git a/packer/_types/__init__.py b/packer/_types/__init__.py
index e45c947..8cc180c 100644
--- a/packer/_types/__init__.py
+++ b/packer/_types/__init__.py
@@ -1,3 +1,4 @@
 from .base import *
 from .float_types import *
 from .int_types import *
+from .sized_packers import *
\ No newline at end of file
diff --git a/packer/_types/sized_packers.py b/packer/_types/sized_packers.py
new file mode 100644
index 0000000..a1583e5
--- /dev/null
+++ b/packer/_types/sized_packers.py
@@ -0,0 +1,32 @@
+__all__ = ("SizedData",)
+
+from typing import Type, Callable
+from .base import TypeDescriptor
+
+def _pack_sized_data_wrapper(size: int) -> Callable:
+    def inner_packer(val: bytes) -> bytes:
+        if (data_len:=len(val)) < size:
+            raise ValueError(f"Attempted to pack data of size {data_len}, expected data of size {size}.")
+        return val
+    return inner_packer
+    
+def _unpack_sized_data_wrapper(size: int) -> Callable:
+    def inner_unpack(data: bytearray) -> tuple[int, bytearray]:
+        return (size, data[:size])
+    return inner_unpack
+
+class SizedDataMeta(type):
+    def __getitem__(cls, data_size: int) -> Type["SizedData"]:
+        return type(
+            f"{cls.__name__}({data_size})", 
+            (TypeDescriptor,), 
+            {
+                "__data_size__": data_size,
+                "pack": _pack_sized_data_wrapper(data_size),
+                "unpack": _unpack_sized_data_wrapper(data_size),
+            }
+        )
+
+class SizedData(metaclass=SizedDataMeta):
+    def pack(val: bytes) -> bytes: ...
+    def unpack(data: bytearray) -> tuple[int, bytearray]: ...
\ No newline at end of file
diff --git a/tests/test_sized_packing.py b/tests/test_sized_packing.py
new file mode 100644
index 0000000..6ea21ad
--- /dev/null
+++ b/tests/test_sized_packing.py
@@ -0,0 +1,31 @@
+import math
+import struct
+from dataclasses import (
+    dataclass,
+)
+
+from packer import (
+    SizedData,
+    Pack,
+    packable,
+)
+
+
+@packable
+@dataclass
+class SizedDataStruct:
+    sized_data: Pack[SizedData[10]] = None
+    sized_data2: Pack[SizedData[50]] = None
+
+def test_sized_packing() -> bool:
+    t = SizedDataStruct()
+    t.sized_data = bytearray(b"h"*5 + b"i"*5)
+    t.sized_data2 = bytearray(50)
+
+    assert len(t.pack()) == 60 and t.pack() == bytearray(b"hhhhhiiiii\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
+    assert t.unpack(bytearray(b"x"*5 + b"d"*5) + bytearray(50)) == 60
+    assert t.sized_data == bytearray(b"x"*5 + b"d"*5)
+    assert t.sized_data2 == bytearray(50)
+
+if __name__ == "__main__":
+    test_sized_packing()