Add read/write access to PacketSideData

Author: skeskinen

av/packet.pxd

@@ -7,11 +7,15 @@ from av.bytesource cimport ByteSource
 from av.stream cimport Stream
 
 
-cdef class PacketSideData:
+cdef class PacketSideData(Buffer):
     cdef uint8_t *data
     cdef size_t size
     cdef lib.AVPacketSideDataType dtype
 
+    cdef size_t _buffer_size(self)
+    cdef void* _buffer_ptr(self)
+    cdef bint _buffer_writable(self)
+
 cdef class Packet(Buffer):
     cdef lib.AVPacket* ptr
     cdef Stream _stream

av/packet.py

@@ -2,6 +2,7 @@
 
 import cython
 from cython.cimports import libav as lib
+from cython.cimports.av.buffer import Buffer
 from cython.cimports.av.bytesource import bytesource
 from cython.cimports.av.error import err_check
 from cython.cimports.av.opaque import opaque_container
@@ -64,7 +65,7 @@ def packet_sidedata_type_from_literal(dtype: PktSideDataT) -> lib.AVPacketSideDa
 
 
 @cython.cclass
-class PacketSideData:
+class PacketSideData(Buffer):
     @staticmethod
     def from_packet(packet: Packet, data_type: PktSideDataT) -> PacketSideData:
         """create new PacketSideData by copying an existing packet's side data
@@ -152,6 +153,19 @@ def data_size(self) -> int:
         """
         return self.size
 
+    # Buffer protocol implementation
+    @cython.cfunc
+    def _buffer_size(self) -> cython.size_t:
+        return self.size
+
+    @cython.cfunc
+    def _buffer_ptr(self) -> cython.p_void:
+        return self.data
+
+    @cython.cfunc
+    def _buffer_writable(self) -> cython.bint:
+        return True
+
     def __bool__(self) -> bool:
         """
         True if this object holds side data.

av/packet.pyi

@@ -50,7 +50,7 @@ PktSideDataT = Literal[
     "rtcp_sr",
 ]
 
-class PacketSideData:
+class PacketSideData(Buffer):
     @staticmethod
     def from_packet(packet: Packet, dtype: PktSideDataT) -> PacketSideData: ...
     def to_packet(self, packet: Packet, move: bool = False): ...

tests/test_packet.py

@@ -1,9 +1,10 @@
+import struct
 from typing import get_args
 from unittest import SkipTest
 
 import av
 
-from .common import fate_suite
+from .common import fate_suite, sandboxed
 
 
 class TestProperties:
@@ -100,3 +101,65 @@ def test_palette(self) -> None:
 
             nxt.set_sidedata(sdata, move=True)
             assert not bool(sdata)
+
+    def test_buffer_protocol(self) -> None:
+        with av.open(fate_suite("h264/extradata-reload-multi-stsd.mov")) as container:
+            for pkt in container.demux():
+                if pkt.has_sidedata("new_extradata"):
+                    sdata = pkt.get_sidedata("new_extradata")
+
+                    raw = bytes(sdata)
+                    assert len(raw) == sdata.data_size > 0
+                    assert sdata.buffer_size == sdata.data_size
+                    assert sdata.buffer_ptr != 0
+                    assert bytes(memoryview(sdata)) == raw
+
+                    # Modify and verify changes stick
+                    sdata.update(b"\xde\xad\xbe\xef" + raw[4:])
+                    assert bytes(sdata)[:4] == b"\xde\xad\xbe\xef"
+
+                    pkt.set_sidedata(sdata)
+                    assert (
+                        bytes(pkt.get_sidedata("new_extradata"))[:4]
+                        == b"\xde\xad\xbe\xef"
+                    )
+                    return
+
+        raise AssertionError("No packet with new_extradata side data found")
+
+    def test_skip_samples_remux(self) -> None:
+        # Source file has skip_end=706 on last packet. Setting to 0 should
+        # result in 706 more decoded samples. And the file duration reported by
+        # the container should also increase.
+        output_path = sandboxed("skip_samples_modified.mkv")
+
+        with av.open(fate_suite("mkv/codec_delay_opus.mkv")) as c:
+            original_samples = sum(f.samples for f in c.decode(c.streams.audio[0]))
+
+        with av.open(fate_suite("mkv/codec_delay_opus.mkv")) as inp:
+            original_duration = inp.duration
+            audio_stream = inp.streams.audio[0]
+            with av.open(output_path, "w") as out:
+                out_stream = out.add_stream_from_template(audio_stream)
+                for pkt in inp.demux(audio_stream):
+                    if pkt.dts is None:
+                        continue
+                    if pkt.has_sidedata("skip_samples"):
+                        sdata = pkt.get_sidedata("skip_samples")
+                        raw = bytes(sdata)
+                        skip_end = struct.unpack("<I", raw[4:8])[0]
+                        assert skip_end == 706
+                        sdata.update(raw[:4] + struct.pack("<I", 0) + raw[8:])
+                        pkt.set_sidedata(sdata)
+                    pkt.stream = out_stream
+                    out.mux(pkt)
+
+        with av.open(output_path) as c:
+            modified_samples = sum(f.samples for f in c.decode(c.streams.audio[0]))
+            modified_duration = c.duration
+
+        assert modified_samples - original_samples == 706
+
+        assert original_duration is not None
+        assert modified_duration is not None
+        assert modified_duration > original_duration