Add type checking

Author: Schamper

dissect/util/_build.py

@@ -1,4 +1,5 @@
 # Reference: https://setuptools.pypa.io/en/latest/build_meta.html#dynamic-build-dependencies-and-other-build-meta-tweaks
+# type: ignore
 from __future__ import annotations
 
 import os

dissect/util/_native/hash/__init__.py dissect/util/_native/hash/__init__.pyidissect/util/_native/hash/__init__.py renamed to dissect/util/_native/hash/__init__.pyi

@@ -23,8 +23,8 @@
 try:
     from dissect.util import _native
 
-    lz4 = lz4_native = _native.compression.lz4
-    lzo = lzo_native = _native.compression.lzo
+    lz4 = lz4_native = _native.compression.lz4  # type: ignore
+    lzo = lzo_native = _native.compression.lzo  # type: ignore
 except (ImportError, AttributeError):
     lz4_native = lzo_native = None
 

dissect/util/_native/hash/crc32c.py dissect/util/_native/hash/crc32c.pyidissect/util/_native/hash/crc32c.py renamed to dissect/util/_native/hash/crc32c.pyi

@@ -25,10 +25,10 @@ def _get_length(src: BinaryIO, length: int) -> int:
 
 
 def decompress(
-    src: bytes | BinaryIO,
+    src: bytes | bytearray | memoryview | BinaryIO,
     uncompressed_size: int = -1,
     return_bytearray: bool = False,
-) -> bytes | tuple[bytes, int]:
+) -> bytes | bytearray | tuple[bytes | bytearray, int]:
     """LZ4 decompress from a file-like object or bytes up to a certain length. Assumes no header.
 
     Args:
@@ -39,7 +39,7 @@ def decompress(
     Returns:
         The decompressed data.
     """
-    if not hasattr(src, "read"):
+    if isinstance(src, bytes | bytearray | memoryview):
         src = io.BytesIO(src)
 
     dst = bytearray()

dissect/util/compression/__init__.py

@@ -11,7 +11,7 @@
 _H = struct.Struct("<H")
 
 
-def decompress(src: bytes | BinaryIO) -> bytes:
+def decompress(src: bytes | bytearray | memoryview | BinaryIO) -> bytes:
     """LZBITMAP decompress from a file-like object or bytes.
 
     Decompresses until EOF or EOS of the input data.
@@ -22,7 +22,7 @@ def decompress(src: bytes | BinaryIO) -> bytes:
     Returns:
         The decompressed data.
     """
-    if not hasattr(src, "read"):
+    if isinstance(src, bytes | bytearray | memoryview):
         src = io.BytesIO(src)
 
     if src.read(4) != b"ZBM\x09":
@@ -54,7 +54,7 @@ def decompress(src: bytes | BinaryIO) -> bytes:
             buf = memoryview(src.read(compressed_size))
 
             # Build the bitmap/token map
-            token_map = []
+            token_map: list[tuple[int | None, int]] = []
             bits = int.from_bytes(buf[-17:], "little")
             for i in range(0xF):
                 if i < 3:
@@ -97,7 +97,7 @@ def decompress(src: bytes | BinaryIO) -> bytes:
 
                 for _ in range(repeat):
                     bitmap, token = token_map[idx]
-                    if idx < 3:
+                    if bitmap is None:  # idx < 3, but this makes the type checker happy
                         # Index 0, 1, 2 are special and indicate we need to read a bitmap from the bitmap region
                         bitmap = buf[bitmap_offset]
                         bitmap_offset += 1

dissect/util/compression/lz4.py

@@ -398,7 +398,7 @@ def _decode_lmd(
     return bytes(dst)
 
 
-def decompress(src: bytes | BinaryIO) -> bytes:
+def decompress(src: bytes | bytearray | memoryview | BinaryIO) -> bytes:
     """LZFSE decompress from a file-like object or bytes.
 
     Decompresses until EOF or EOS of the input data.
@@ -409,7 +409,7 @@ def decompress(src: bytes | BinaryIO) -> bytes:
     Returns:
         The decompressed data.
     """
-    if not hasattr(src, "read"):
+    if isinstance(src, bytes | bytearray | memoryview):
         src = io.BytesIO(src)
 
     dst = bytearray()

dissect/util/compression/lzbitmap.py

@@ -25,7 +25,7 @@ def _get_displacement(offset: int) -> int:
 TAG_MASKS = [(1 << i) for i in range(8)]
 
 
-def decompress(src: bytes | BinaryIO) -> bytes:
+def decompress(src: bytes | bytearray | memoryview | BinaryIO) -> bytes:
     """LZNT1 decompress from a file-like object or bytes.
 
     Args:
@@ -34,7 +34,7 @@ def decompress(src: bytes | BinaryIO) -> bytes:
     Returns:
         The decompressed data.
     """
-    if not hasattr(src, "read"):
+    if isinstance(src, bytes | bytearray | memoryview):
         src = io.BytesIO(src)
 
     offset = src.tell()

dissect/util/compression/lzfse.py

@@ -23,7 +23,7 @@ def _read_length(src: BinaryIO, val: int, mask: int) -> int:
     return length + mask + val
 
 
-def decompress(src: bytes | BinaryIO, header: bool = True, buflen: int = -1) -> bytes:
+def decompress(src: bytes | bytearray | memoryview | BinaryIO, header: bool = True, buflen: int = -1) -> bytes:
     """LZO decompress from a file-like object or bytes. Assumes no header.
 
     Arguments are largely compatible with python-lzo API.
@@ -36,7 +36,7 @@ def decompress(src: bytes | BinaryIO, header: bool = True, buflen: int = -1) ->
     Returns:
         The decompressed data.
     """
-    if not hasattr(src, "read"):
+    if isinstance(src, bytes | bytearray | memoryview):
         src = io.BytesIO(src)
 
     dst = bytearray()

dissect/util/compression/lznt1.py

@@ -56,7 +56,7 @@
 _H = struct.Struct("<H")
 
 
-def decompress(src: bytes | BinaryIO) -> bytes:
+def decompress(src: bytes | bytearray | memoryview | BinaryIO) -> bytes:
     """LZVN decompress from a file-like object or bytes.
 
     Decompresses until EOF or EOS of the input data.
@@ -67,7 +67,7 @@ def decompress(src: bytes | BinaryIO) -> bytes:
     Returns:
         The decompressed data.
     """
-    if not hasattr(src, "read"):
+    if isinstance(src, bytes | bytearray | memoryview):
         src = io.BytesIO(src)
 
     offset = src.tell()
@@ -207,7 +207,7 @@ def decompress(src: bytes | BinaryIO) -> bytes:
             if src_size < opc_len:
                 break
 
-            src_size -= opc_len + L
+            src_size -= opc_len
             break
 
         elif opc in OP_UDEF: