Add float helpers (#167)

Author: Schamper

dissect/cstruct/__init__.py

@@ -30,12 +30,18 @@
 )
 from dissect.cstruct.util import (
     dumpstruct,
+    f16,
+    f32,
+    f64,
     hexdump,
     p8,
     p16,
     p32,
     p64,
     pack,
+    pf16,
+    pf32,
+    pf64,
     swap,
     swap16,
     swap32,
@@ -75,12 +81,18 @@
     "ctypes",
     "ctypes_type",
     "dumpstruct",
+    "f16",
+    "f32",
+    "f64",
     "hexdump",
     "p8",
     "p16",
     "p32",
     "p64",
     "pack",
+    "pf16",
+    "pf32",
+    "pf64",
     "swap",
     "swap16",
     "swap32",

dissect/cstruct/compiler.py

@@ -11,6 +11,7 @@
 from dissect.cstruct.bitbuffer import BitBuffer
 from dissect.cstruct.types import (
     Array,
+    BaseArray,
     Char,
     CharArray,
     Flag,
@@ -24,18 +25,15 @@
     Wchar,
     WcharArray,
 )
-from dissect.cstruct.types.base import BaseArray
 from dissect.cstruct.types.enum import EnumMetaType
-from dissect.cstruct.types.packed import _struct
+from dissect.cstruct.util import _struct
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
     from types import MethodType
 
     from dissect.cstruct.cstruct import cstruct
-    from dissect.cstruct.types import (
-        BaseType,
-    )
+    from dissect.cstruct.types import BaseType
     from dissect.cstruct.types.structure import Field
 
 SUPPORTED_TYPES = (

dissect/cstruct/types/base.py

@@ -6,6 +6,7 @@
 
 from dissect.cstruct.exception import ArraySizeError
 from dissect.cstruct.expression import Expression
+from dissect.cstruct.util import normalize_endianness
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -328,27 +329,5 @@ def _is_eof(stream: BinaryIO) -> bool:
     return False
 
 
-ENDIANNESS_MAP: dict[AllowedEndianness, Endianness] = {
-    "<": "<",
-    ">": ">",
-    "!": "!",
-    "@": "@",
-    "=": "=",
-    "network": "!",
-    "little": "<",
-    "big": ">",
-}
-
-
-def normalize_endianness(endian: AllowedEndianness) -> Endianness:
-    """Normalize an endianness string to one of the standard format characters."""
-    try:
-        return ENDIANNESS_MAP[endian]
-    except KeyError:
-        raise ValueError(
-            f"Invalid endianness: {endian!r}, expected one of {', '.join(ENDIANNESS_MAP.keys())}"
-        ) from None
-
-
 # As mentioned in the BaseType class, we correctly set the type here
 MetaType.ArrayType = Array

dissect/cstruct/types/packed.py

@@ -1,22 +1,16 @@
 from __future__ import annotations
 
-from functools import lru_cache
-from struct import Struct
 from typing import TYPE_CHECKING, Any, BinaryIO, Generic, TypeVar
 
 from dissect.cstruct.types.base import EOF, BaseType
+from dissect.cstruct.util import _struct
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
     from dissect.cstruct.cstruct import Endianness
 
 
-@lru_cache(1024)
-def _struct(endian: str, packchar: str) -> Struct:
-    return Struct(f"{endian}{packchar}")
-
-
 T = TypeVar("T", int, float)
 
 

dissect/cstruct/util.py

@@ -3,15 +3,17 @@
 import os
 import pprint
 import string
+import struct
 import sys
 from enum import Enum
+from functools import lru_cache
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
     from typing import Literal
 
-    from dissect.cstruct.cstruct import AllowedEndianness
+    from dissect.cstruct.cstruct import AllowedEndianness, Endianness
     from dissect.cstruct.types.base import BaseType
     from dissect.cstruct.types.structure import Structure
 
@@ -64,6 +66,28 @@
     "big": "big",
 }
 
+ENDIANNESS_MAP: dict[AllowedEndianness, Endianness] = {
+    "<": "<",
+    ">": ">",
+    "!": "!",
+    "@": "@",
+    "=": "=",
+    "network": "!",
+    "little": "<",
+    "big": ">",
+}
+
+
+def normalize_endianness(endian: AllowedEndianness) -> Endianness:
+    """Normalize an endianness string to one of the standard format characters."""
+    try:
+        return ENDIANNESS_MAP[endian]
+    except KeyError:
+        raise ValueError(
+            f"Invalid endianness: {endian!r}, expected one of {', '.join(ENDIANNESS_MAP.keys())}"
+        ) from None
+
+
 Palette = list[tuple[int, str]]
 
 
@@ -327,6 +351,11 @@ def dumpstruct(
     return _dumpstruct(obj, obj.dumps(), offset, color, output, autoskip)
 
 
+@lru_cache(1024)
+def _struct(endian: str, packchar: str) -> struct.Struct:
+    return struct.Struct(f"{endian}{packchar}")
+
+
 def pack(value: int, size: int | None = None, endian: AllowedEndianness = "little") -> bytes:
     """Pack an integer value to a given bit size, endianness.
 
@@ -400,6 +429,36 @@ def p64(value: int, endian: AllowedEndianness = "little") -> bytes:
     return pack(value, 64, endian)
 
 
+def pf16(value: float, endian: AllowedEndianness = "little") -> bytes:
+    """Pack a 16 bit float.
+
+    Arguments:
+        value: Value to pack.
+        endian: Endianness to use (little, big, network, <, >, !, @ or =).
+    """
+    return _struct(normalize_endianness(endian), "e").pack(value)
+
+
+def pf32(value: float, endian: AllowedEndianness = "little") -> bytes:
+    """Pack a 32 bit float.
+
+    Arguments:
+        value: Value to pack.
+        endian: Endianness to use (little, big, network, <, >, !, @ or =).
+    """
+    return _struct(normalize_endianness(endian), "f").pack(value)
+
+
+def pf64(value: float, endian: AllowedEndianness = "little") -> bytes:
+    """Pack a 64 bit float.
+
+    Arguments:
+        value: Value to pack.
+        endian: Endianness to use (little, big, network, <, >, !, @ or =).
+    """
+    return _struct(normalize_endianness(endian), "d").pack(value)
+
+
 def u8(value: bytes, endian: AllowedEndianness = "little", sign: bool = False) -> int:
     """Unpack an 8 bit integer.
 
@@ -444,6 +503,36 @@ def u64(value: bytes, endian: AllowedEndianness = "little", sign: bool = False)
     return unpack(value, 64, endian, sign)
 
 
+def f16(value: bytes, endian: AllowedEndianness = "little") -> float:
+    """Unpack a 16 bit float.
+
+    Arguments:
+        value: Value to unpack.
+        endian: Endianness to use (little, big, network, <, >, !, @ or =).
+    """
+    return _struct(normalize_endianness(endian), "e").unpack(value)[0]
+
+
+def f32(value: bytes, endian: AllowedEndianness = "little") -> float:
+    """Unpack a 32 bit float.
+
+    Arguments:
+        value: Value to unpack.
+        endian: Endianness to use (little, big, network, <, >, !, @ or =).
+    """
+    return _struct(normalize_endianness(endian), "f").unpack(value)[0]
+
+
+def f64(value: bytes, endian: AllowedEndianness = "little") -> float:
+    """Unpack a 64 bit float.
+
+    Arguments:
+        value: Value to unpack.
+        endian: Endianness to use (little, big, network, <, >, !, @ or =).
+    """
+    return _struct(normalize_endianness(endian), "d").unpack(value)[0]
+
+
 def swap(value: int, size: int) -> int:
     """Swap the endianness of an integer with a given bit size.
 

tests/test_utils.py

@@ -259,54 +259,84 @@ def test_pack_unpack() -> None:
     assert util.p16(1, endian) == b"\x01\x00"
     assert util.p32(1, endian) == b"\x01\x00\x00\x00"
     assert util.p64(1, endian) == b"\x01\x00\x00\x00\x00\x00\x00\x00"
+    assert util.pf16(1.0, endian) == b"\x00\x3c"
+    assert util.pf32(1.0, endian) == b"\x00\x00\x80\x3f"
+    assert util.pf64(1.0, endian) == b"\x00\x00\x00\x00\x00\x00\xf0\x3f"
     assert util.u8(b"\x01", endian, sign) == 1
     assert util.u16(b"\x01\x00", endian, sign) == 1
     assert util.u32(b"\x01\x00\x00\x00", endian, sign) == 1
     assert util.u64(b"\x01\x00\x00\x00\x00\x00\x00\x00", endian, sign) == 1
+    assert util.f16(b"\x00\x3c", endian) == 1.0
+    assert util.f32(b"\x00\x00\x80\x3f", endian) == 1.0
+    assert util.f64(b"\x00\x00\x00\x00\x00\x00\xf0\x3f", endian) == 1.0
 
     endian = "big"
     sign = False
     assert util.p8(1, endian) == b"\x01"
     assert util.p16(1, endian) == b"\x00\x01"
     assert util.p32(1, endian) == b"\x00\x00\x00\x01"
     assert util.p64(1, endian) == b"\x00\x00\x00\x00\x00\x00\x00\x01"
+    assert util.pf16(1.0, endian) == b"\x3c\x00"
+    assert util.pf32(1.0, endian) == b"\x3f\x80\x00\x00"
+    assert util.pf64(1.0, endian) == b"\x3f\xf0\x00\x00\x00\x00\x00\x00"
     assert util.u8(b"\x01", endian, sign) == 1
     assert util.u16(b"\x00\x01", endian, sign) == 1
     assert util.u32(b"\x00\x00\x00\x01", endian, sign) == 1
     assert util.u64(b"\x00\x00\x00\x00\x00\x00\x00\x01", endian, sign) == 1
+    assert util.f16(b"\x3c\x00", endian) == 1.0
+    assert util.f32(b"\x3f\x80\x00\x00", endian) == 1.0
+    assert util.f64(b"\x3f\xf0\x00\x00\x00\x00\x00\x00", endian) == 1.0
 
     endian = "network"
     sign = False
     assert util.p8(1, endian) == b"\x01"
     assert util.p16(1, endian) == b"\x00\x01"
     assert util.p32(1, endian) == b"\x00\x00\x00\x01"
     assert util.p64(1, endian) == b"\x00\x00\x00\x00\x00\x00\x00\x01"
+    assert util.pf16(1.0, endian) == b"\x3c\x00"
+    assert util.pf32(1.0, endian) == b"\x3f\x80\x00\x00"
+    assert util.pf64(1.0, endian) == b"\x3f\xf0\x00\x00\x00\x00\x00\x00"
     assert util.u8(b"\x01", endian, sign) == 1
     assert util.u16(b"\x00\x01", endian, sign) == 1
     assert util.u32(b"\x00\x00\x00\x01", endian, sign) == 1
     assert util.u64(b"\x00\x00\x00\x00\x00\x00\x00\x01", endian, sign) == 1
+    assert util.f16(b"\x3c\x00", endian) == 1.0
+    assert util.f32(b"\x3f\x80\x00\x00", endian) == 1.0
+    assert util.f64(b"\x3f\xf0\x00\x00\x00\x00\x00\x00", endian) == 1.0
 
     endian = "little"
     sign = True
     assert util.p8(-120, endian) == b"\x88"
     assert util.p16(-120, endian) == b"\x88\xff"
     assert util.p32(-120, endian) == b"\x88\xff\xff\xff"
     assert util.p64(-120, endian) == b"\x88\xff\xff\xff\xff\xff\xff\xff"
+    assert util.pf16(-120.0, endian) == b"\x80\xd7"
+    assert util.pf32(-120.0, endian) == b"\x00\x00\xf0\xc2"
+    assert util.pf64(-120.0, endian) == b"\x00\x00\x00\x00\x00\x00\x5e\xc0"
     assert util.u8(b"\x88", endian, sign) == -120
     assert util.u16(b"\x88\xff", endian, sign) == -120
     assert util.u32(b"\x88\xff\xff\xff", endian, sign) == -120
     assert util.u64(b"\x88\xff\xff\xff\xff\xff\xff\xff", endian, sign) == -120
+    assert util.f16(b"\x80\xd7", endian) == -120.0
+    assert util.f32(b"\x00\x00\xf0\xc2", endian) == -120.0
+    assert util.f64(b"\x00\x00\x00\x00\x00\x00\x5e\xc0", endian) == -120.0
 
     endian = "big"
     sign = True
     assert util.p8(-120, endian) == b"\x88"
     assert util.p16(-120, endian) == b"\xff\x88"
     assert util.p32(-120, endian) == b"\xff\xff\xff\x88"
     assert util.p64(-120, endian) == b"\xff\xff\xff\xff\xff\xff\xff\x88"
+    assert util.pf16(-120.0, endian) == b"\xd7\x80"
+    assert util.pf32(-120.0, endian) == b"\xc2\xf0\x00\x00"
+    assert util.pf64(-120.0, endian) == b"\xc0\x5e\x00\x00\x00\x00\x00\x00"
     assert util.u8(b"\x88", endian, sign) == -120
     assert util.u16(b"\xff\x88", endian, sign) == -120
     assert util.u32(b"\xff\xff\xff\x88", endian, sign) == -120
     assert util.u64(b"\xff\xff\xff\xff\xff\xff\xff\x88", endian, sign) == -120
+    assert util.f16(b"\xd7\x80", endian) == -120.0
+    assert util.f32(b"\xc2\xf0\x00\x00", endian) == -120.0
+    assert util.f64(b"\xc0\x5e\x00\x00\x00\x00\x00\x00", endian) == -120.0
 
     assert util.pack(1, 24) == b"\x01\x00\x00"
     assert util.unpack(b"\x01\x00\x00", 24) == 1