Widen type requirement for unpack to allow buffer protocol objects

dataclasses_struct/_typing.py

@@ -10,8 +10,14 @@
 else:
     from typing_extensions import Unpack, dataclass_transform
 
+if sys.version_info >= (3, 12):
+    from collections.abc import Buffer
+else:
+    from typing_extensions import Buffer
+
 
 __all__ = [
+    "Buffer",
     "TypeGuard",
     "Unpack",
     "dataclass_transform",

dataclasses_struct/dataclass.py

@@ -19,7 +19,7 @@
     overload,
 )
 
-from ._typing import TypeGuard, Unpack, dataclass_transform
+from ._typing import Buffer, TypeGuard, Unpack, dataclass_transform
 from .field import Field, builtin_fields
 from .types import PadAfter, PadBefore
 
@@ -192,7 +192,7 @@ def _init_from_args(self, args: Iterator) -> T:
             setattr(obj, name, arg)
         return obj
 
-    def _unpack(self, data: bytes) -> T:
+    def _unpack(self, data: Buffer) -> T:
         return self._init_from_args(iter(self.struct.unpack(data)))
 
 
@@ -206,7 +206,7 @@ class DataclassStructProtocol(Protocol):
     """
 
     @classmethod
-    def from_packed(cls: type[T], data: bytes) -> T:
+    def from_packed(cls: type[T], data: Buffer) -> T:
         """Return an instance of the class from its packed representation.
 
         Args:
@@ -529,12 +529,12 @@ def pack(self) -> bytes:
 
 def _make_unpack_method(cls: type) -> classmethod:
     func = """
-def from_packed(cls, data: bytes) -> cls_type:
+def from_packed(cls, data: Buffer) -> cls_type:
     '''Unpack from bytes.'''
     return cls.__dataclass_struct__._unpack(data)
 """
 
-    scope: dict[str, Any] = {"cls_type": cls}
+    scope: dict[str, Any] = {"cls_type": cls, "Buffer": Buffer}
     exec(func, {}, scope)
     return classmethod(scope["from_packed"])
 

dataclasses_struct/ext/mypy_plugin.py

@@ -11,6 +11,7 @@
 
 
 def transform_dataclass_struct(ctx: ClassDefContext) -> bool:
+    buffer_type = ctx.api.named_type("dataclasses_struct._typing.Buffer")
     bytes_type = ctx.api.named_type("builtins.bytes")
     tvd = TypeVarType(
         "T",
@@ -25,7 +26,14 @@ def transform_dataclass_struct(ctx: ClassDefContext) -> bool:
         ctx.api,
         ctx.cls,
         "from_packed",
-        [Argument(Var("data", bytes_type), bytes_type, None, ArgKind.ARG_POS)],
+        [
+            Argument(
+                Var("data", buffer_type),
+                buffer_type,
+                None,
+                ArgKind.ARG_POS,
+            )
+        ],
         tvd,
         self_type=TypeType(tvd),
         tvar_def=tvd,

docs/guide.md

@@ -30,18 +30,30 @@ controls the size and alignment of fields:
 
 Decorated classes are transformed to a standard Python
 [dataclass](https://docs.python.org/3/library/dataclasses.html) with boilerplate
-`__init__`, `__repr__`, `__eq__` etc. auto-generated. Additionally, two methods
-are added to the class:
-[`pack`][dataclasses_struct.DataclassStructProtocol.pack], a method for packing
-an instance of the class to `bytes`, and
-[`from_packed`][dataclasses_struct.DataclassStructProtocol.from_packed], a class
-method that returns a new instance of the class from its packed `bytes`
-representation. The additional `dataclass_kwargs` keyword arguments will be
-passed through to the [stdlib `dataclass`
-decorator](https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass):
+`__init__`, `__repr__`, `__eq__` etc. auto-generated. The additional
+`dataclass_kwargs` keyword arguments will be passed through to the [stdlib
+`dataclass` decorator](https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass):
 all standard keyword arguments are supported except for `slots` and
 `weakref_slot`.
 
+
+In addition to the standard `dataclass` methods, two methods
+are added to the class:
+
+* [`pack`][dataclasses_struct.DataclassStructProtocol.pack], which packs an
+  instance of the class to `bytes`.
+* [`from_packed`][dataclasses_struct.DataclassStructProtocol.from_packed], which
+  is a class method that returns a new instance of the class from its packed
+  representation in an object that implements the [buffer
+  protococol](https://docs.python.org/3/c-api/buffer.html) (`bytes`,
+  `bytearray`, [memory-mapped file
+  objects](https://docs.python.org/3/library/mmap.html) etc.).
+
+A class attribute named
+[`__dataclass_struct__`][dataclasses_struct.DataclassStructProtocol.__dataclass_struct__]
+is also added (see [Inspecting
+dataclass-structs](#inspecting-dataclass-structs)).
+
 ## Default value validation
 
 Default attribute values will be validated against their expected type and
@@ -492,6 +504,7 @@ can add stubs for the generated functions and attribute to the class:
 
 ```python
 from typing import ClassVar, TYPE_CHECKING
+from collections.abc import Buffer  # import from typing_extensions on Python <3.12
 import dataclasses_struct as dcs
 
 @dcs.dataclass_struct()
@@ -505,7 +518,7 @@ class Test:
         def pack(self) -> bytes: ...
 
         @classmethod
-        def from_packed(cls, data: bytes) -> "Test": ...
+        def from_packed(cls, data: Buffer) -> "Test": ...
 ```
 
 The [`DataclassStructProtocol`][dataclasses_struct.DataclassStructProtocol]

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 ]
 requires-python = ">=3.9.0"
 dependencies = [
-    "typing-extensions>=4.12.2 ; python_full_version < '3.11'",
+    "typing-extensions>=4.12.2 ; python_full_version < '3.12'",
 ]
 
 [project.urls]

test/test-mypy-plugin.yml

@@ -50,3 +50,37 @@
 
     t = Test.from_packed(Test(1).pack())
     reveal_type(t)  # N: Revealed type is "main.Test"
+
+- case: test_from_packed_supports_bytearray_argument
+  mypy_config: 'plugins = dataclasses_struct.ext.mypy_plugin'
+  main: |
+    import dataclasses_struct as dcs
+
+    @dcs.dataclass_struct()
+    class Test:
+        x: int
+
+    t = Test.from_packed(bytearray(Test(1).pack()))
+    reveal_type(t)  # N: Revealed type is "main.Test"
+
+- case: test_from_packed_supports_mmap_argument
+  mypy_config: 'plugins = dataclasses_struct.ext.mypy_plugin'
+  main: |
+    import mmap
+    import tempfile
+    from pathlib import Path
+
+    import dataclasses_struct as dcs
+
+    @dcs.dataclass_struct()
+    class Test:
+        x: int
+
+    packed = Test(1).pack()
+    with tempfile.TemporaryDirectory() as tempdir:
+        path = Path(tempdir) / "data"
+        path.write_bytes(packed)
+        with path.open("rb+") as f, mmap.mmap(f.fileno(), 0) as mapped:
+          t = Test.from_packed(mapped)
+
+    reveal_type(t)  # N: Revealed type is "main.Test"

test/test_pack_unpack.py

@@ -1,6 +1,8 @@
 import dataclasses
 import itertools
+import mmap
 import struct
+from pathlib import Path
 from typing import Annotated
 
 import pytest
@@ -648,3 +650,37 @@ class T:
     t.x = -1
     unpacked = T.from_packed(t.pack())
     assert unpacked.x == -1
+
+
+def test_pack_unpack_bytearray() -> None:
+    @dcs.dataclass_struct()
+    class T:
+        x: dcs.Int = 100
+        y: dcs.F32 = -1.5
+        z: Annotated[bytes, 12] = b"hello, world"
+
+    t = T()
+    packed_array = bytearray(t.pack())
+    assert T.from_packed(packed_array) == t
+
+
+def test_pack_unpack_mmap(tmp_path: Path) -> None:
+    @dcs.dataclass_struct()
+    class T:
+        x: dcs.Int = 100
+        y: dcs.F32 = -1.5
+        z: Annotated[bytes, 12] = b"hello, world"
+
+    path = tmp_path / "data"
+
+    t = T()
+    packed = t.pack()
+    path.write_bytes(packed)
+
+    with path.open("rb+") as f, mmap.mmap(f.fileno(), 0) as mapped:
+        unpacked = T.from_packed(mapped)
+
+    assert unpacked == t
+
+    # Check that the file isn't changed
+    assert path.read_bytes() == packed