added packet side-data handler

Author: tikuma-lsuhsc

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

av/packet.py

@@ -4,6 +4,194 @@
 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
+
+# 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_extra_data",
+    "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_name: PacketSideDataTypeLiteral,
+) -> lib.AVPacketSideDataType:
+
+    i = cython.declare(
+        cython.int, get_args(PacketSideDataTypeLiteral).index(data_type_name)
+    )
+
+    return cython.cast(lib.AVPacketSideDataType, i)
+
+
+def packet_side_data_type_desc(data_type_name: PacketSideDataTypeLiteral) -> str | None:
+    """FFmpeg description of packet side data type"""
+    dtype = _packet_side_data_type_from_literal(data_type_name)
+    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)
+        c_ptr = lib.av_packet_side_data_get(
+            packet.ptr.side_data, packet.ptr.side_data_elems, dtype
+        )
+
+        found = cython.declare(cython.bint, c_ptr != cython.NULL)
+
+        sdata = PacketSideData(dtype, c_ptr.size if found else 0)
+
+        with cython.nogil:
+            if found:
+                memcpy(sdata.data, c_ptr.data, c_ptr.size)
+
+        return sdata
+
+    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.cclass
@@ -235,3 +423,40 @@ 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)

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

test1/rawvideo_pal8.avi

@@ -11,7 +11,7 @@
 
 import numpy as np
 
-from av.datasets import fate as fate_suite
+from av.datasets import fate as fate_suite, curated as pyav_curated
 
 try:
     import PIL  # noqa
@@ -28,7 +28,7 @@
     T = TypeVar("T")
 
 
-__all__ = ("fate_suite",)
+__all__ = ("fate_suite", "pyav_curated")
 
 
 is_windows = os.name == "nt"

tests/common.py

@@ -1,6 +1,8 @@
 import av
 
-from .common import fate_suite
+from .common import fate_suite, pyav_curated
+
+from typing import get_args
 
 
 class TestProperties:
@@ -48,3 +50,41 @@ 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)
+        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) != None
+
+    def test_palette(self):
+
+        with av.open(pyav_curated("ffmpeg/rawvideo_pal8.avi")) as container:
+            # its first packet has palette side data
+
+            iterpackets = container.demux()
+            pkt = next(iterpackets)  # has palette
+
+            assert pkt.has_side_data("palette")
+
+            sdata = pkt.get_side_data("palette")
+            assert sdata.data_type == "palette"
+            assert bool(sdata)
+            assert sdata.data_size == 1024
+            assert sdata.data_desc == 'Palette'
+
+            nxt = next(iterpackets)  # has no palette
+
+            assert not nxt.has_side_data("palette")
+
+            sdata1 = nxt.get_side_data("palette")
+            assert sdata1.data_type == "palette"
+            assert not bool(sdata1)
+            assert sdata1.data_size == 0
+
+            nxt.set_side_data(sdata, move=True)
+            assert not bool(sdata)