diff --git a/av/packet.pxd b/av/packet.pxd index f1517a3d4..8e8ad034d 100644 --- a/av/packet.pxd +++ b/av/packet.pxd @@ -1,3 +1,5 @@ +from cython.cimports.libc.stdint import uint8_t + cimport libav as lib from av.buffer cimport Buffer @@ -5,6 +7,11 @@ from av.bytesource cimport ByteSource from av.stream cimport Stream +cdef class PacketSideData: + cdef uint8_t *data + cdef size_t size + cdef lib.AVPacketSideDataType dtype + cdef class Packet(Buffer): cdef lib.AVPacket* ptr cdef Stream _stream diff --git a/av/packet.py b/av/packet.py index 81f1aaa4d..e1051b653 100644 --- a/av/packet.py +++ b/av/packet.py @@ -1,9 +1,183 @@ +from typing import Iterator, Literal, get_args + import cython from cython.cimports import libav as lib from cython.cimports.av.bytesource import bytesource 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.libc.string import memcpy + +# 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 +PktSideDataT = 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_sidedata_type_to_literal(dtype: lib.AVPacketSideDataType) -> PktSideDataT: + return get_args(PktSideDataT)[cython.cast(int, dtype)] + + +def packet_sidedata_type_from_literal(dtype: PktSideDataT) -> lib.AVPacketSideDataType: + return get_args(PktSideDataT).index(dtype) + + +@cython.cclass +class PacketSideData: + @staticmethod + def from_packet(packet: Packet, data_type: PktSideDataT) -> 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_sidedata_type_from_literal(data_type) + return _packet_sidedata_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_sidedata_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_sidedata_from_packet( + packet: cython.pointer[lib.AVPacket], dtype: lib.AVPacketSideDataType +) -> PacketSideData: + with cython.nogil: + c_ptr = lib.av_packet_side_data_get( + packet.side_data, packet.side_data_elems, dtype + ) + found: 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 @cython.cclass @@ -235,3 +409,48 @@ def opaque(self, v): if v is None: return self.ptr.opaque_ref = opaque_container.add(v) + + def has_sidedata(self, dtype: str) -> bool: + """True if this packet has the specified side data + + :param dtype: side data type + :type dtype: str + """ + + dtype2 = packet_sidedata_type_from_literal(dtype) + return ( + lib.av_packet_side_data_get( + self.ptr.side_data, self.ptr.side_data_elems, dtype2 + ) + != cython.NULL + ) + + def get_sidedata(self, dtype: str) -> PacketSideData: + """get a copy of the side data + + :param dtype: side data type (:method:`~av.packet.PacketSideData.sidedata_types` for the full list of options) + :type dtype: 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, dtype) + + def set_sidedata(self, sidedata: PacketSideData, move: cython.bint = False): + """copy or move side data to this packet + + :param sidedata: Source packet side data + :type sidedata: :class:`~av.packet.PacketSideData` + :param move: If True, move the data from `sidedata` object, defaults to False + :type move: bool + """ + sidedata.to_packet(self, move) + + def iter_sidedata(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_sidedata_from_packet(self.ptr, self.ptr.side_data[i].type) diff --git a/av/packet.pyi b/av/packet.pyi index baa234d7b..8dab20065 100644 --- a/av/packet.pyi +++ b/av/packet.pyi @@ -1,10 +1,70 @@ from fractions import Fraction +from typing import Iterator, Literal from av.subtitles.subtitle import SubtitleSet from .buffer import Buffer from .stream import Stream +# Sync with definition in 'packet.py' +PktSideDataT = 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", +] + +class PacketSideData: + @staticmethod + def from_packet(packet: Packet, dtype: PktSideDataT) -> PacketSideData: ... + def to_packet(self, packet: Packet, move: bool = False): ... + @property + def data_type(self) -> str: ... + @property + def data_desc(self) -> str: ... + @property + def data_size(self) -> int: ... + def __bool__(self) -> bool: ... + +def packet_sidedata_type_to_literal(dtype: int) -> PktSideDataT: ... +def packet_sidedata_type_from_literal(dtype: PktSideDataT) -> int: ... + class Packet(Buffer): stream: Stream stream_index: int @@ -23,3 +83,7 @@ class Packet(Buffer): def __init__(self, input: int | bytes | None = None) -> None: ... def decode(self) -> list[SubtitleSet]: ... + def has_sidedata(self, dtype: PktSideDataT) -> bool: ... + def get_sidedata(self, dtype: PktSideDataT) -> PacketSideData: ... + def set_sidedata(self, sidedata: PacketSideData, move: bool = False) -> None: ... + def iter_sidedata(self) -> Iterator[PacketSideData]: ... diff --git a/include/libavcodec/avcodec.pxd b/include/libavcodec/avcodec.pxd index 2bce8c2c1..43ecb0260 100644 --- 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 index 423396f71..d5552238c 100644 --- a/tests/test_packet.py +++ b/tests/test_packet.py @@ -1,3 +1,5 @@ +from typing import get_args + import av from .common import fate_suite @@ -48,3 +50,46 @@ def test_set_duration(self) -> None: packet.duration += 10 assert packet.duration == old_duration + 10 + + +class TestPacketSideData: + def test_data_types(self) -> None: + dtypes = get_args(av.packet.PktSideDataT) + ffmpeg_ver = [int(v) for v in av.ffmpeg_version_info.split(".", 2)[:2]] + for dtype in dtypes: + av_enum = av.packet.packet_sidedata_type_from_literal(dtype) + assert dtype == av.packet.packet_sidedata_type_to_literal(av_enum) + + if (ffmpeg_ver[0] < 8 and dtype == "lcevc") or ( + ffmpeg_ver[0] < 9 and dtype == "rtcp_sr" + ): + break + + def test_iter(self) -> None: + with av.open(fate_suite("h264/extradata-reload-multi-stsd.mov")) as container: + for pkt in container.demux(): + for sdata in pkt.iter_sidedata(): + assert pkt.dts == 2 and sdata.data_type == "new_extradata" + + def test_palette(self) -> None: + 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_sidedata("new_extradata")) + + sdata = pkt.get_sidedata("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_sidedata("new_extradata") + + sdata1 = nxt.get_sidedata("new_extradata") + assert sdata1.data_type == "new_extradata" + assert not bool(sdata1) + assert sdata1.data_size == 0 + + nxt.set_sidedata(sdata, move=True) + assert not bool(sdata)