Coverage for dataclasses_struct/dataclass.py: 97%
187 statements
« prev ^ index » next coverage.py v7.8.2, created at 2025-06-10 09:31 +1200
1import dataclasses
2import sys
3from collections.abc import Generator, Iterator
4from struct import Struct
5from typing import (
6 Annotated,
7 Any,
8 Callable,
9 Generic,
10 Literal,
11 Protocol,
12 TypedDict,
13 TypeVar,
14 Union,
15 get_args,
16 get_origin,
17 get_type_hints,
18 overload,
19)
21from ._typing import TypeGuard, Unpack, dataclass_transform
22from .field import Field, builtin_fields
23from .types import PadAfter, PadBefore
26def _get_padding_and_field(fields):
27 pad_before = pad_after = 0
28 field = None
29 for f in fields:
30 if isinstance(f, PadBefore):
31 pad_before += f.size
32 elif isinstance(f, PadAfter):
33 pad_after += f.size
34 elif field is not None:
35 raise TypeError(f"too many annotations: {f}")
36 else:
37 field = f
39 return pad_before, pad_after, field
42T = TypeVar("T")
45_SIZE_BYTEORDER_MODE_CHAR: dict[tuple[str, str], str] = {
46 ("native", "native"): "@",
47 ("std", "native"): "=",
48 ("std", "little"): "<",
49 ("std", "big"): ">",
50 ("std", "network"): "!",
51}
52_MODE_CHAR_SIZE_BYTEORDER: dict[str, tuple[str, str]] = {
53 v: k for k, v in _SIZE_BYTEORDER_MODE_CHAR.items()
54}
57class _DataclassStructInternal(Generic[T]):
58 struct: Struct
59 cls: type[T]
60 _fieldnames: list[str]
61 _fieldtypes: list[type]
63 @property
64 def format(self) -> str:
65 return self.struct.format
67 @property
68 def size(self) -> int:
69 return self.struct.size
71 @property
72 def mode(self) -> str:
73 return self.format[0]
75 def __init__(
76 self,
77 fmt: str,
78 cls: type,
79 fieldnames: list[str],
80 fieldtypes: list[type],
81 ):
82 self.struct = Struct(fmt)
83 self.cls = cls
84 self._fieldnames = fieldnames
85 self._fieldtypes = fieldtypes
87 def _flattened_attrs(self, obj) -> list[Any]:
88 """
89 Returns a list of all attributes, including those of any nested structs
90 """
91 attrs = []
92 for fieldname in self._fieldnames:
93 attr = getattr(obj, fieldname)
94 if is_dataclass_struct(attr):
95 attrs.extend(attr.__dataclass_struct__._flattened_attrs(attr))
96 else:
97 attrs.append(attr)
98 return attrs
100 def pack(self, obj: T) -> bytes:
101 return self.struct.pack(*self._flattened_attrs(obj))
103 def _arg_generator(self, args: Iterator) -> Generator:
104 for fieldtype in self._fieldtypes:
105 if is_dataclass_struct(fieldtype):
106 yield fieldtype.__dataclass_struct__._init_from_args(args)
107 else:
108 yield fieldtype(next(args))
110 def _init_from_args(self, args: Iterator) -> T:
111 """
112 Returns an instance of self.cls, consuming args
113 """
114 return self.cls(*self._arg_generator(args))
116 def unpack(self, data: bytes) -> T:
117 return self._init_from_args(iter(self.struct.unpack(data)))
120class DataclassStructProtocol(Protocol):
121 __dataclass_struct__: _DataclassStructInternal
123 @classmethod
124 def from_packed(cls: type[T], data: bytes) -> T: ...
126 def pack(self) -> bytes: ...
129@overload
130def is_dataclass_struct(
131 obj: type,
132) -> TypeGuard[type[DataclassStructProtocol]]: ...
135@overload
136def is_dataclass_struct(obj: object) -> TypeGuard[DataclassStructProtocol]: ...
139def is_dataclass_struct(
140 obj: Union[type, object],
141) -> Union[
142 TypeGuard[DataclassStructProtocol],
143 TypeGuard[type[DataclassStructProtocol]],
144]:
145 """
146 Returns True if obj is a class that has been decorated with
147 dataclasses_struct.dataclass or an instance of one.
148 """
149 return (
150 dataclasses.is_dataclass(obj)
151 and hasattr(obj, "__dataclass_struct__")
152 and isinstance(obj.__dataclass_struct__, _DataclassStructInternal)
153 )
156def get_struct_size(cls_or_obj) -> int:
157 """
158 Returns the size of the packed representation of the struct in bytes.
159 Accepts either a class or an instance of a dataclass_struct.
160 """
161 if not is_dataclass_struct(cls_or_obj):
162 raise TypeError(f"{cls_or_obj} is not a dataclass_struct")
163 return cls_or_obj.__dataclass_struct__.size
166class _BytesField(Field[bytes]):
167 field_type = bytes
169 def __init__(self, n: int):
170 if not isinstance(n, int) or n < 1:
171 raise ValueError("bytes length must be positive non-zero int")
173 self.n = n
175 def format(self) -> str:
176 return f"{self.n}s"
178 def validate_default(self, val: bytes) -> None:
179 if len(val) > self.n:
180 raise ValueError(f"bytes cannot be longer than {self.n} bytes")
182 def __repr__(self) -> str:
183 return f"{super().__repr__()}({self.n})"
186class _NestedField(Field):
187 field_type: type[DataclassStructProtocol]
189 def __init__(self, cls: type[DataclassStructProtocol]):
190 self.field_type = cls
192 def format(self) -> str:
193 # Return the format without the byteorder specifier at the beginning
194 return self.field_type.__dataclass_struct__.format[1:]
197def _validate_and_parse_field(
198 cls: type,
199 name: str,
200 field_type: type,
201 is_native: bool,
202 validate_defaults: bool,
203 mode: str,
204) -> tuple[str, type]:
205 """
206 name is the name of the attribute, f is its type annotation.
207 """
209 if get_origin(field_type) == Annotated:
210 # The types defined in .types (e.g. U32, F32, etc.) are of the form:
211 # Annotated[<builtin type>, Field(<field args>)]
212 # Alternatively, accept annotations of the form:
213 # Annotated[<builtin type>, PadBefore(<size>), PadAfter(<size>)]
214 # or:
215 # Annotated[<dcs.types type>, PadBefore(<size>), PadAfter(<size>)]
216 # or:
217 # Annotated[bytes, <positive non-zero integer>]
218 # which will expand to:
219 # Annotated[
220 # <builtin type associated with dcs.types type>,
221 # Field(<field args>),
222 # PadBefore(<size>),
223 # PadAfter(<size>),
224 # ]
225 # PadBefore and PadAfter are optional and may be repeated or in
226 # different orders
227 type_, *fields = get_args(field_type)
228 pad_before, pad_after, field = _get_padding_and_field(fields)
229 else:
230 # Accept type annotations without Annotated e.g. builtin types or
231 # nested dataclass-structs
232 pad_before = pad_after = 0
233 type_ = field_type
234 field = None
236 if field is None:
237 # Must be either a nested type or one of the supported builtins
238 if is_dataclass_struct(type_):
239 nested_mode = type_.__dataclass_struct__.mode
240 if nested_mode != mode:
241 size, byteorder = _MODE_CHAR_SIZE_BYTEORDER[nested_mode]
242 exp_size, exp_byteorder = _MODE_CHAR_SIZE_BYTEORDER[mode]
243 msg = (
244 "byteorder and size of nested dataclass-struct does not "
245 f"match that of container (expected '{exp_size}' size and "
246 f"'{exp_byteorder}' byteorder, got '{size}' size and "
247 f"'{byteorder}' byteorder)"
248 )
249 raise TypeError(msg)
250 field = _NestedField(type_)
251 else:
252 field = builtin_fields.get(type_)
253 if field is None:
254 raise TypeError(f"type not supported: {field_type}")
256 if not isinstance(field, Field):
257 if issubclass(type_, bytes):
258 # Annotated[bytes, <positive non-zero integer>] is a byte array
259 field = _BytesField(field)
260 else:
261 raise TypeError(f"invalid field annotation: {field!r}")
262 elif not issubclass(type_, field.field_type):
263 raise TypeError(f"type {type_} not supported for field: {field}")
265 if is_native:
266 if not field.is_native:
267 raise TypeError(
268 f"field {field} only supported in standard size mode"
269 )
270 elif not field.is_std:
271 raise TypeError(f"field {field} only supported in native size mode")
273 if validate_defaults and hasattr(cls, name):
274 val = getattr(cls, name)
275 if not isinstance(val, field.field_type):
276 raise TypeError(
277 "invalid type for field: expected "
278 f"{field.field_type} got {type(val)}"
279 )
280 field.validate_default(val)
282 return (
283 "".join(
284 (
285 (f"{pad_before}x" if pad_before else ""),
286 field.format(),
287 (f"{pad_after}x" if pad_after else ""),
288 )
289 ),
290 type_,
291 )
294def _make_pack_method() -> Callable:
295 func = """
296def pack(self) -> bytes:
297 '''Pack to bytes using struct.pack.'''
298 return self.__dataclass_struct__.pack(self)
299"""
301 scope: dict[str, Any] = {}
302 exec(func, {}, scope)
303 return scope["pack"]
306def _make_unpack_method(cls: type) -> classmethod:
307 func = """
308def from_packed(cls, data: bytes) -> cls_type:
309 '''Unpack from bytes.'''
310 return cls.__dataclass_struct__.unpack(data)
311"""
313 scope: dict[str, Any] = {"cls_type": cls}
314 exec(func, {}, scope)
315 return classmethod(scope["from_packed"])
318def _make_class(
319 cls: type,
320 mode: str,
321 is_native: bool,
322 validate_defaults: bool,
323 dataclass_kwargs,
324) -> type[DataclassStructProtocol]:
325 cls_annotations = get_type_hints(cls, include_extras=True)
326 struct_format = [mode]
327 fieldtypes = []
328 for name, field in cls_annotations.items():
329 fmt, type_ = _validate_and_parse_field(
330 cls,
331 name=name,
332 field_type=field,
333 is_native=is_native,
334 validate_defaults=validate_defaults,
335 mode=mode,
336 )
337 struct_format.append(fmt)
338 fieldtypes.append(type_)
340 setattr( # noqa: B010
341 cls,
342 "__dataclass_struct__",
343 _DataclassStructInternal(
344 "".join(struct_format),
345 cls,
346 list(cls_annotations.keys()),
347 fieldtypes,
348 ),
349 )
350 setattr(cls, "pack", _make_pack_method()) # noqa: B010
351 setattr(cls, "from_packed", _make_unpack_method(cls)) # noqa: B010
353 return dataclasses.dataclass(cls, **dataclass_kwargs)
356class _DataclassKwargsPre310(TypedDict, total=False):
357 init: bool
358 repr: bool
359 eq: bool
360 order: bool
361 unsafe_hash: bool
362 frozen: bool
365if sys.version_info >= (3, 10):
367 class DataclassKwargs(_DataclassKwargsPre310, total=False):
368 match_args: bool
369 kw_only: bool
370else:
372 class DataclassKwargs(_DataclassKwargsPre310, total=False):
373 pass
376@overload
377def dataclass_struct(
378 *,
379 size: Literal["native"] = "native",
380 byteorder: Literal["native"] = "native",
381 validate_defaults: bool = True,
382 **dataclass_kwargs: Unpack[DataclassKwargs],
383) -> Callable[[type], type]: ...
386@overload
387def dataclass_struct(
388 *,
389 size: Literal["std"],
390 byteorder: Literal["native", "big", "little", "network"] = "native",
391 validate_defaults: bool = True,
392 **dataclass_kwargs: Unpack[DataclassKwargs],
393) -> Callable[[type], type]: ...
396@dataclass_transform()
397def dataclass_struct(
398 *,
399 size: Literal["native", "std"] = "native",
400 byteorder: Literal["native", "big", "little", "network"] = "native",
401 validate_defaults: bool = True,
402 **dataclass_kwargs: Unpack[DataclassKwargs],
403) -> Callable[[type], type]:
404 is_native = size == "native"
405 if is_native:
406 if byteorder != "native":
407 raise ValueError("'native' size requires 'native' byteorder")
408 elif size != "std":
409 raise ValueError(f"invalid size: {size}")
410 if byteorder not in ("native", "big", "little", "network"):
411 raise ValueError(f"invalid byteorder: {byteorder}")
413 for kwarg in ("slots", "weakref_slot"):
414 if dataclass_kwargs.get(kwarg):
415 msg = f"dataclass '{kwarg}' keyword argument is not supported"
416 raise ValueError(msg)
418 def decorator(cls: type) -> type:
419 return _make_class(
420 cls,
421 mode=_SIZE_BYTEORDER_MODE_CHAR[(size, byteorder)],
422 is_native=is_native,
423 validate_defaults=validate_defaults,
424 dataclass_kwargs=dataclass_kwargs,
425 )
427 return decorator