added packet side-data handler

tikuma-lsuhsc · tikuma-lsuhsc · commit a965bdaaba1d · 2025-09-01T20:53:11.000-04:00
diff --git a/av/packet.pxd b/av/packet.pxd
@@ -4,6 +4,12 @@ from av.buffer cimport Buffer
 from av.bytesource cimport ByteSource
 from av.stream cimport Stream
 
+from cython.cimports.libc.stdint import uint8_t
+
+cdef class PacketSideData:
+    cdef uint8_t *data
+    cdef size_t size
+    cdef lib.AVPacketSideDataType dtype
 
 cdef class Packet(Buffer):
     cdef lib.AVPacket* ptr
@@ -13,3 +19,7 @@ cdef class Packet(Buffer):
     cdef ByteSource source
     cdef size_t _buffer_size(self)
     cdef void* _buffer_ptr(self)
+
+cdef PacketSideData _packet_side_data_from_packet(
+    lib.AVPacket* packet, lib.AVPacketSideDataType data_type
+)
diff --git a/av/packet.py b/av/packet.py
@@ -4,6 +4,203 @@
 from cython.cimports.av.error import err_check
 from cython.cimports.av.opaque import opaque_container
 from cython.cimports.av.utils import avrational_to_fraction, to_avrational
+from cython.cimports.av.error import err_check
+from cython.cimports.libc.string import memcpy
+
+from typing import Literal, get_args, Iterator
+
+# check https://github.com/FFmpeg/FFmpeg/blob/master/libavcodec/packet.h#L41 for
+# new additions in the future FFmpeg releases
+# note: the order must follow that of the AVPacketSideDataType enum def
+PacketSideDataTypeLiteral = Literal[
+    "palette",
+    "new_extradata",
+    "param_change",
+    "h263_mb_info",
+    "replay_gain",
+    "display_matrix",
+    "stereo_3d",
+    "audio_service_type",
+    "quality_stats",
+    "fallback_track",
+    "cpb_properties",
+    "skip_samples",
+    "jp_dual_mono",
+    "strings_metadata",
+    "subtitle_position",
+    "matroska_block_additional",
+    "webvtt_identifier",
+    "webvtt_settings",
+    "metadata_update",
+    "mpegts_stream_id",
+    "mastering_display_metadata",
+    "spherical",
+    "content_light_level",
+    "a53_cc",
+    "encryption_init_info",
+    "encryption_info",
+    "afd",
+    "prft",
+    "icc_profile",
+    "dovi_conf",
+    "s12m_timecode",
+    "dynamic_hdr10_plus",
+    "iamf_mix_gain_param",
+    "iamf_info_param",
+    "iamf_recon_gain_info_param",
+    "ambient_viewing_environment",
+    "frame_cropping",
+    "lcevc",
+    "3d_reference_displays",
+    "rtcp_sr",
+]
+
+
+def _packet_side_data_type_to_literal(
+    data_type: lib.AVPacketSideDataType,
+) -> PacketSideDataTypeLiteral:
+
+    return get_args(PacketSideDataTypeLiteral)[cython.cast(int, data_type)]
+
+
+def _packet_side_data_type_from_literal(
+    data_type: PacketSideDataTypeLiteral,
+) -> lib.AVPacketSideDataType:
+
+    i = cython.declare(
+        cython.int, get_args(PacketSideDataTypeLiteral).index(data_type)
+    )
+
+    return cython.cast(lib.AVPacketSideDataType, i)
+
+
+def packet_side_data_type_desc(data_type: PacketSideDataTypeLiteral) -> str | None:
+    """FFmpeg description of packet side data type"""
+    dtype = _packet_side_data_type_from_literal(data_type)
+    res = lib.av_packet_side_data_name(dtype)
+    return res if res != cython.NULL else None
+
+
+@cython.cclass
+class PacketSideData:
+
+    @staticmethod
+    def from_packet(
+        packet: "Packet", data_type: PacketSideDataTypeLiteral
+    ) -> "PacketSideData":
+        """create new PacketSideData by copying an existing packet's side data
+
+        :param packet: Source packet
+        :type packet: :class:`~av.packet.Packet`
+        :param data_type: side data type
+        :return: newly created copy of the side data if the side data of the
+                 requested type is found in the packet, else an empty object
+        :rtype: :class:`~av.packet.PacketSideData`
+        """
+
+        dtype = _packet_side_data_type_from_literal(data_type)
+        return _packet_side_data_from_packet(packet.ptr, dtype)
+
+    def __cinit__(self, dtype: lib.AVPacketSideDataType, size: cython.size_t):
+        self.dtype = dtype
+        with cython.nogil:
+            if size:
+                self.data = cython.cast(cython.p_uchar, lib.av_malloc(size))
+                if self.data == cython.NULL:
+                    raise MemoryError("Failed to allocate memory")
+            else:
+                self.data = cython.NULL
+        self.size = size
+
+    def __dealloc__(self):
+        with cython.nogil:
+            lib.av_freep(cython.address(self.data))
+
+    def to_packet(self, packet: "Packet", move: cython.bint = False):
+        """copy or move side data to the specified packet
+
+        :param packet: Target packet
+        :type packet: :class:`~av.packet.Packet`
+        :param move: True to move the data from this object to the packet,
+                     defaults to False.
+        :type move: bool
+        """
+        if self.size == 0:
+            # nothing to add, should clear existing side_data in packet?
+            return
+
+        data = self.data
+
+        with cython.nogil:
+            if not move:
+                data = cython.cast(cython.p_uchar, lib.av_malloc(self.size))
+                if data == cython.NULL:
+                    raise MemoryError("Failed to allocate memory")
+                memcpy(data, self.data, self.size)
+
+            res = lib.av_packet_add_side_data(packet.ptr, self.dtype, data, self.size)
+        err_check(res)
+
+        if move:
+            self.data = cython.NULL
+            self.size = 0
+
+    @property
+    def data_type(self) -> str:
+        """
+        The type of this packet side data.
+
+        :type: str
+        """
+        return _packet_side_data_type_to_literal(self.dtype)
+
+    @property
+    def data_desc(self) -> str:
+        """
+        The description of this packet side data type.
+
+        :type: str
+        """
+
+        return lib.av_packet_side_data_name(self.dtype)
+
+    @property
+    def data_size(self) -> int:
+        """
+        The size in bytes of this packet side data.
+
+        :type: int
+        """
+        return self.size
+
+    def __bool__(self) -> bool:
+        """
+        True if this object holds side data.
+
+        :type: bool
+        """
+        return self.data != cython.NULL
+
+
+@cython.cfunc
+def _packet_side_data_from_packet(
+    packet: cython.pointer[lib.AVPacket], data_type: lib.AVPacketSideDataType
+) -> PacketSideData:
+
+    with cython.nogil:
+        c_ptr = lib.av_packet_side_data_get(
+            packet.side_data, packet.side_data_elems, data_type
+        )
+
+        found = cython.declare(cython.bint, c_ptr != cython.NULL)
+
+    sdata = PacketSideData(data_type, c_ptr.size if found else 0)
+
+    with cython.nogil:
+        if found:
+            memcpy(sdata.data, c_ptr.data, c_ptr.size)
+
+    return sdata
 
 
 @cython.cclass
@@ -235,3 +432,49 @@ def opaque(self, v):
         if v is None:
             return
         self.ptr.opaque_ref = opaque_container.add(v)
+
+    def has_side_data(self, data_type: str) -> bool:
+        """True if this packet has the specified side data
+
+        :param data_type: side data type
+        :type data_type: str
+        """
+
+        dtype = _packet_side_data_type_from_literal(data_type)
+        return (
+            lib.av_packet_side_data_get(
+                self.ptr.side_data, self.ptr.side_data_elems, dtype
+            )
+            != cython.NULL
+        )
+
+    def get_side_data(self, data_type: str) -> PacketSideData:
+        """get a copy of the side data
+
+        :param data_type: side data type (:method:`~av.packet.PacketSideData.side_data_types` for the full list of options)
+        :type data_type: str
+        :return: newly created copy of the side data if the side data of the
+                 requested type is found in the packet, else an empty object
+        :rtype: :class:`~av.packet.PacketSideData`
+        """
+        return PacketSideData.from_packet(self, data_type)
+
+    def set_side_data(self, side_data: PacketSideData, move: bool = False):
+        """copy or move side data to this packet
+
+        :param side_data: Source packet side data
+        :type side_data: :class:`~av.packet.PacketSideData`
+        :param move: True to move the data from `side_data` object,
+                     defaults to False.
+        :type move: bool
+        """
+        side_data.to_packet(self, move)
+
+    def iter_side_data(self) -> Iterator[PacketSideData]:
+        """iterate over side data of this packet.
+
+        :yield: :class:`~av.packet.PacketSideData` object
+        """
+
+        for i in range(self.ptr.side_data_elems):
+            yield _packet_side_data_from_packet(self.ptr, self.ptr.side_data[i].type)
diff --git a/include/libavcodec/avcodec.pxd b/include/libavcodec/avcodec.pxd
@@ -1,4 +1,4 @@
-from libc.stdint cimport int8_t, int64_t, uint16_t, uint32_t
+from libc.stdint cimport int8_t, int64_t, uint16_t, uint32_t, uint8_t
 
 cdef extern from "libavcodec/codec.h":
     struct AVCodecTag:
@@ -17,6 +17,17 @@ cdef extern from "libavcodec/packet.h" nogil:
         int free_opaque
     )
 
+    const AVPacketSideData *av_packet_side_data_get(const AVPacketSideData *sd,
+                                                int nb_sd,
+                                                AVPacketSideDataType type)
+
+    uint8_t* av_packet_get_side_data(const AVPacket *pkt, AVPacketSideDataType type,
+                                    size_t *size)
+
+    int av_packet_add_side_data(AVPacket *pkt, AVPacketSideDataType type,
+                                uint8_t *data, size_t size)
+
+    const char *av_packet_side_data_name(AVPacketSideDataType type)
 
 cdef extern from "libavutil/channel_layout.h":
     ctypedef enum AVChannelOrder:
@@ -469,6 +480,8 @@ cdef extern from "libavcodec/avcodec.h" nogil:
         int size
         int stream_index
         int flags
+        AVPacketSideData *side_data
+        int side_data_elems
         int duration
         int64_t pos
         void *opaque
diff --git a/tests/test_packet.py b/tests/test_packet.py
@@ -2,6 +2,8 @@
 
 from .common import fate_suite
 
+from typing import get_args
+
 
 class TestProperties:
     def test_is_keyframe(self) -> None:
@@ -48,3 +50,50 @@ def test_set_duration(self) -> None:
                 packet.duration += 10
 
                 assert packet.duration == old_duration + 10
+
+
+class TestPacketSideData:
+
+    def test_data_types(self):
+        dtypes = get_args(av.packet.PacketSideDataTypeLiteral)
+        ffmpeg_ver = [int(v) for v in av.ffmpeg_version_info.split(".", 2)[:2]]
+        for dtype in dtypes:
+            av_enum = av.packet._packet_side_data_type_from_literal(dtype)
+            assert dtype == av.packet._packet_side_data_type_to_literal(av_enum)
+            assert av.packet.packet_side_data_type_desc(dtype) is not None
+
+            if (ffmpeg_ver[0] < 8 and dtype == "lcevc") or (
+                ffmpeg_ver[0] < 9 and dtype == "rtcp_sr"
+            ):
+                break
+
+    def test_iter(self):
+        with av.open(fate_suite("h264/extradata-reload-multi-stsd.mov")) as container:
+            for pkt in container.demux():
+                for sdata in pkt.iter_side_data():
+                    assert pkt.dts == 2 and sdata.data_type == "new_extradata"
+
+    def test_palette(self):
+
+        with av.open(fate_suite("h264/extradata-reload-multi-stsd.mov")) as container:
+
+            iterpackets = container.demux()
+            pkt = next(pkt for pkt in iterpackets if pkt.has_side_data("new_extradata"))
+
+            sdata = pkt.get_side_data("new_extradata")
+            assert sdata.data_type == "new_extradata"
+            assert bool(sdata)
+            assert sdata.data_size > 0
+            assert sdata.data_desc == "New Extradata"
+
+            nxt = next(iterpackets)  # has no palette
+
+            assert not nxt.has_side_data("new_extradata")
+
+            sdata1 = nxt.get_side_data("new_extradata")
+            assert sdata1.data_type == "new_extradata"
+            assert not bool(sdata1)
+            assert sdata1.data_size == 0
+
+            nxt.set_side_data(sdata, move=True)
+            assert not bool(sdata)
