Add type-hints

Author: Schamper

dissect/util/_build.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import setuptools.build_meta
+from setuptools.build_meta import *  # noqa: F403
+
+
+def get_requires_for_build_wheel(config_settings: dict[str, Any] | None = None) -> list[str]:
+    base = setuptools.build_meta.get_requires_for_build_wheel(config_settings)
+
+    if os.getenv("BUILD_RUST", "").lower() in {"1", "true"} or (config_settings and "--build-rust" in config_settings):
+        return [*base, "setuptools-rust"]
+
+    return base

dissect/util/_native.src/src/compression/lz4.rs

@@ -4,6 +4,19 @@ use pyo3::types::{PyByteArray, PyBytes};
 
 const MAX_DISCOVER_OUTPUT_SIZE: usize = 1024 * 1024 * 1024;
 
+/// LZ4 decompress bytes up to a certain length. Assumes no header.
+///
+/// Unlike the Python implementation, this function does not support streaming decompression
+/// (i.e. reading from a file-like object).
+///
+/// Args:
+///     src: Bytes to decompress.
+///     uncompressed_size: The uncompressed data size. If not provided or ``-1``, will try to discover it.
+///     return_bytearray: Whether to return ``bytearray`` or ``bytes``.
+///
+/// Returns:
+///     The decompressed data or a tuple of the decompressed data and the amount of bytes read.
+///
 #[pyfunction]
 #[pyo3(signature = (src, uncompressed_size=-1, return_bytearray=false))]
 fn decompress(

dissect/util/_native.src/src/compression/lzo.rs

@@ -2,6 +2,20 @@ use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 
+/// LZO decompress from bytes. Assumes no header.
+///
+/// Arguments are largely compatible with python-lzo API.
+///
+/// Unlike the Python implementation, this function does not support streaming decompression
+/// (i.e. reading from a file-like object).
+/// Args:
+///     src: Bytes to decompress.
+///     header: Whether the metadata header is included in the input.
+///     buflen: If ``header`` is ``False``, a buffer length in bytes must be given that will fit the output.
+///
+/// Returns:
+///     The decompressed data.
+///
 #[pyfunction]
 #[pyo3(signature = (src, header=true, buflen=-1))]
 fn decompress(

dissect/util/_native.src/src/lib.rs

@@ -2,7 +2,6 @@ use pyo3::prelude::*;
 
 mod compression;
 
-
 #[pymodule]
 fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
     compression::create_submodule(m)?;

dissect/util/_native/__init__.pyi

@@ -0,0 +1,3 @@
+from dissect.util._native import compression
+
+__all__ = ["compression"]

dissect/util/_native/compression/__init__.pyi

@@ -0,0 +1,3 @@
+from dissect.util._native.compression import lz4, lzo
+
+__all__ = ["lz4", "lzo"]

dissect/util/_native/compression/lz4.pyi

@@ -0,0 +1,7 @@
+from typing import BinaryIO
+
+def decompress(
+    src: bytes | BinaryIO,
+    uncompressed_size: int = -1,
+    return_bytearray: bool = False,
+) -> bytes | tuple[bytes, int]: ...

dissect/util/_native/compression/lzo.pyi

@@ -0,0 +1,3 @@
+from typing import BinaryIO
+
+def decompress(src: bytes | BinaryIO, header: bool = True, buflen: int = -1) -> bytes: ...

dissect/util/compression/__init__.py

@@ -25,7 +25,7 @@
 
     lz4 = lz4_native = _native.compression.lz4
     lzo = lzo_native = _native.compression.lzo
-except ImportError:
+except (ImportError, AttributeError):
     lz4_native = lzo_native = None
 
 __all__ = [

dissect/util/compression/lz4.py

@@ -29,10 +29,10 @@ def decompress(
     uncompressed_size: int = -1,
     return_bytearray: bool = False,
 ) -> bytes | tuple[bytes, int]:
-    """LZ4 decompress from a file-like object up to a certain length. Assumes no header.
+    """LZ4 decompress from a file-like object or bytes up to a certain length. Assumes no header.
 
     Args:
-        src: File-like object to decompress from.
+        src: File-like object or bytes to decompress from.
         uncompressed_size: Ignored, present for compatibility with native lz4.
         return_bytearray: Whether to return ``bytearray`` or ``bytes``.