frameparams: support defining frame format for get_frame_parameters()

Author: swehner

c-ext/decompressor.c

@@ -281,6 +281,7 @@ PyObject *Decompressor_decompress(ZstdDecompressor *self, PyObject *args,
     size_t zresult;
     ZSTD_outBuffer outBuffer;
     ZSTD_inBuffer inBuffer;
+    ZSTD_frameHeader frameHeader;
 
     if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|nOO:decompress", kwlist,
                                      &source, &maxOutputSize, &readAcrossFrames,
@@ -299,15 +300,14 @@ PyObject *Decompressor_decompress(ZstdDecompressor *self, PyObject *args,
         goto finally;
     }
 
-    decompressedSize = ZSTD_getFrameContentSize(source.buf, source.len);
-
-    if (ZSTD_CONTENTSIZE_ERROR == decompressedSize) {
+    if (ZSTD_getFrameHeader_advanced(&frameHeader, source.buf, source.len, self->format) != 0) {
         PyErr_SetString(ZstdError,
                         "error determining content size from frame header");
         goto finally;
     }
+    decompressedSize = frameHeader.frameContentSize;
     /* Special case of empty frame. */
-    else if (0 == decompressedSize) {
+    if (0 == decompressedSize) {
         result = PyBytes_FromStringAndSize("", 0);
         goto finally;
     }

c-ext/frameparams.c

@@ -12,19 +12,20 @@ extern PyObject *ZstdError;
 
 FrameParametersObject *get_frame_parameters(PyObject *self, PyObject *args,
                                             PyObject *kwargs) {
-    static char *kwlist[] = {"data", NULL};
+    static char *kwlist[] = {"data", "format", NULL};
 
     Py_buffer source;
     ZSTD_frameHeader header;
+    ZSTD_format_e format = ZSTD_f_zstd1;
     FrameParametersObject *result = NULL;
     size_t zresult;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:get_frame_parameters",
-                                     kwlist, &source)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|I:get_frame_parameters",
+                                     kwlist, &source, &format)) {
         return NULL;
     }
 
-    zresult = ZSTD_getFrameHeader(&header, source.buf, source.len);
+    zresult = ZSTD_getFrameHeader_advanced(&header, source.buf, source.len, format);
 
     if (ZSTD_isError(zresult)) {
         PyErr_Format(ZstdError, "cannot get frame parameters: %s",

docs/news.rst

@@ -70,6 +70,10 @@ Changes
   an apparant conda bug.
 * Added CI coverage for Anaconda 3.12 and 3.13 on Linux and Windows x86-64
   (but the Windows builds no-op due to above issue).
+* ``get_frame_parameters()`` now accepts an optional ``format`` argument
+  defining the zstandard frame type. You can pass e.g.
+  ``zstandard.FORMAT_ZSTD1_MAGICLESS`` to decode frames without header magic.
+  (#217)
 
 Backwards Compatibility Notes
 -----------------------------

rust-ext/src/decompressor.rs

@@ -19,10 +19,7 @@ use {
         types::{PyBytes, PyList},
         wrap_pyfunction,
     },
-    std::{
-        ffi::c_ulonglong,
-        sync::Arc
-    },
+    std::{ffi::c_ulonglong, sync::Arc},
 };
 
 #[pyclass(module = "zstandard.backend_rust")]
@@ -181,27 +178,35 @@ impl ZstdDecompressor {
 
         self.setup_dctx(py, true)?;
 
-        let output_size =
-            unsafe { zstd_sys::ZSTD_getFrameContentSize(buffer.buf_ptr(), buffer.len_bytes()) };
+        let mut header: zstd_sys::ZSTD_FrameHeader = unsafe { std::mem::zeroed() };
+        let zresult = unsafe {
+            zstd_sys::ZSTD_getFrameHeader_advanced(
+                &mut header,
+                buffer.buf_ptr(),
+                buffer.len_bytes(),
+                self.format,
+            )
+        };
+
+        if zresult != 0 {
+            return Err(ZstdError::new_err(
+                "error determining content size from frame header",
+            ));
+        }
 
-        let (output_buffer_size, output_size) =
-            if output_size == zstd_sys::ZSTD_CONTENTSIZE_ERROR as c_ulonglong {
+        let (output_buffer_size, output_size) = if header.frameContentSize == 0 {
+            return Ok(PyBytes::new(py, &[]));
+        } else if header.frameContentSize == zstd_sys::ZSTD_CONTENTSIZE_UNKNOWN as c_ulonglong {
+            if max_output_size == 0 {
                 return Err(ZstdError::new_err(
-                    "error determining content size from frame header",
+                    "could not determine content size in frame header",
                 ));
-            } else if output_size == 0 {
-                return Ok(PyBytes::new(py, &[]));
-            } else if output_size == zstd_sys::ZSTD_CONTENTSIZE_UNKNOWN as c_ulonglong {
-                if max_output_size == 0 {
-                    return Err(ZstdError::new_err(
-                        "could not determine content size in frame header",
-                    ));
-                }
+            }
 
-                (max_output_size, 0)
-            } else {
-                (output_size as _, output_size)
-            };
+            (max_output_size, 0)
+        } else {
+            (header.frameContentSize as _, header.frameContentSize)
+        };
 
         let mut dest_buffer: Vec<u8> = Vec::new();
         dest_buffer

rust-ext/src/frame_parameters.rs

@@ -6,7 +6,7 @@
 
 use {
     crate::ZstdError,
-    pyo3::{buffer::PyBuffer, prelude::*, wrap_pyfunction},
+    pyo3::{buffer::PyBuffer, exceptions::PyValueError, prelude::*, wrap_pyfunction},
     std::ffi::c_ulonglong,
 };
 
@@ -70,10 +70,22 @@ fn frame_header_size(data: PyBuffer<u8>) -> PyResult<usize> {
 }
 
 #[pyfunction]
-fn get_frame_parameters(py: Python, buffer: PyBuffer<u8>) -> PyResult<Py<FrameParameters>> {
+#[pyo3(signature = (buffer, format=zstd_sys::ZSTD_format_e::ZSTD_f_zstd1 as u32))]
+fn get_frame_parameters(
+    py: Python,
+    buffer: PyBuffer<u8>,
+    format: u32,
+) -> PyResult<Py<FrameParameters>> {
     let raw_data = unsafe {
         std::slice::from_raw_parts::<u8>(buffer.buf_ptr() as *const _, buffer.len_bytes())
     };
+    let format = if format == zstd_sys::ZSTD_format_e::ZSTD_f_zstd1 as u32 {
+        zstd_sys::ZSTD_format_e::ZSTD_f_zstd1
+    } else if format == zstd_sys::ZSTD_format_e::ZSTD_f_zstd1_magicless as u32 {
+        zstd_sys::ZSTD_format_e::ZSTD_f_zstd1_magicless
+    } else {
+        return Err(PyValueError::new_err(format!("invalid format value")));
+    };
 
     let mut header = zstd_sys::ZSTD_FrameHeader {
         frameContentSize: 0,
@@ -87,7 +99,12 @@ fn get_frame_parameters(py: Python, buffer: PyBuffer<u8>) -> PyResult<Py<FramePa
         _reserved2: 0,
     };
     let zresult = unsafe {
-        zstd_sys::ZSTD_getFrameHeader(&mut header, raw_data.as_ptr() as *const _, raw_data.len())
+        zstd_sys::ZSTD_getFrameHeader_advanced(
+            &mut header,
+            raw_data.as_ptr() as *const _,
+            raw_data.len(),
+            format,
+        )
     };
 
     if unsafe { zstd_sys::ZSTD_isError(zresult) } != 0 {

tests/test_compressor_fuzzing.py

@@ -767,7 +767,7 @@ def test_data_equivalence(self, original, threads, use_dict):
         dctx = zstd.ZstdDecompressor(**kwargs)
 
         for i, frame in enumerate(result):
-            self.assertEqual(dctx.decompress(frame), original[i])
+            self.assertEqual(dctx.decompress(frame.tobytes()), original[i])
 
 
 @unittest.skipUnless("ZSTD_SLOW_TESTS" in os.environ, "ZSTD_SLOW_TESTS not set")

tests/test_data_structures.py

@@ -190,6 +190,16 @@ def test_attributes(self):
         self.assertEqual(params.dict_id, 15)
         self.assertTrue(params.has_checksum)
 
+    def test_attributes_without_header(self):
+        # Set multiple things without magic header
+        params = zstd.get_frame_parameters(
+            b"\x45\x40\x0f\x10\x00", format=zstd.FORMAT_ZSTD1_MAGICLESS
+        )
+        self.assertEqual(params.content_size, 272)
+        self.assertEqual(params.window_size, 262144)
+        self.assertEqual(params.dict_id, 15)
+        self.assertTrue(params.has_checksum)
+
     def test_input_types(self):
         v = zstd.FRAME_HEADER + b"\x00\x00"
 

tests/test_decompressor_decompress.py

@@ -37,6 +37,16 @@ def test_input_types(self):
         for source in sources:
             self.assertEqual(dctx.decompress(source), b"foo")
 
+    def test_headerless(self):
+        compression_params = zstd.ZstdCompressionParameters(
+            format=zstd.FORMAT_ZSTD1_MAGICLESS,
+        )
+        cctx = zstd.ZstdCompressor(compression_params=compression_params)
+        compressed = cctx.compress(b"foo")
+
+        dctx = zstd.ZstdDecompressor(format=zstd.FORMAT_ZSTD1_MAGICLESS)
+        self.assertEqual(dctx.decompress(compressed), b"foo")
+
     def test_no_content_size_in_frame(self):
         cctx = zstd.ZstdCompressor(write_content_size=False)
         compressed = cctx.compress(b"foobar")

zstandard/__init__.pyi

@@ -450,7 +450,7 @@ class FrameParameters(object):
 def estimate_decompression_context_size() -> int: ...
 def frame_content_size(data: ByteString) -> int: ...
 def frame_header_size(data: ByteString) -> int: ...
-def get_frame_parameters(data: ByteString) -> FrameParameters: ...
+def get_frame_parameters(data: ByteString, format: Optional[int] = None) -> FrameParameters: ...
 def train_dictionary(
     dict_size: int,
     samples: list[ByteString],

zstandard/backend_cffi.py

@@ -2559,7 +2559,7 @@ def frame_header_size(data):
     return zresult
 
 
-def get_frame_parameters(data):
+def get_frame_parameters(data, format=FORMAT_ZSTD1):
     """
     Parse a zstd frame header into frame parameters.
 
@@ -2570,13 +2570,15 @@ def get_frame_parameters(data):
 
     :param data:
        Data from which to read frame parameters.
+    :param format:
+       Set the format of data for the decoder.
     :return:
        :py:class:`FrameParameters`
     """
     params = ffi.new("ZSTD_FrameHeader *")
 
     data_buffer = ffi.from_buffer(data)
-    zresult = lib.ZSTD_getFrameHeader(params, data_buffer, len(data_buffer))
+    zresult = lib.ZSTD_getFrameHeader_advanced(params, data_buffer, len(data_buffer), format)
     if lib.ZSTD_isError(zresult):
         raise ZstdError(
             "cannot get frame parameters: %s" % _zstd_error(zresult)
@@ -3821,13 +3823,12 @@ def decompress(
 
         data_buffer = ffi.from_buffer(data)
 
-        output_size = lib.ZSTD_getFrameContentSize(
-            data_buffer, len(data_buffer)
-        )
-
-        if output_size == lib.ZSTD_CONTENTSIZE_ERROR:
+        params = ffi.new("ZSTD_frameHeader *")
+        zresult = lib.ZSTD_getFrameHeader_advanced(params, data_buffer, len(data_buffer), self._format)
+        if zresult != 0:
             raise ZstdError("error determining content size from frame header")
-        elif output_size == 0:
+        output_size = params.frameContentSize
+        if output_size == 0:
             return b""
         elif output_size == lib.ZSTD_CONTENTSIZE_UNKNOWN:
             if not max_output_size: