From 1041719934938ece1231aeba8b52217dc93a972b Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Thu, 20 Aug 2026 22:26:15 -0400 Subject: [PATCH] Expose more AVCodecContext fields --- CHANGELOG.rst | 3 + av/audio/codeccontext.py | 11 ++ av/audio/codeccontext.pyi | 2 + av/codec/context.pxd | 1 + av/codec/context.py | 212 +++++++++++++++++++++++++++++++++++- av/codec/context.pyi | 32 +++++- av/video/codeccontext.py | 46 ++++++++ av/video/codeccontext.pyi | 3 + include/avcodec.pxd | 16 +++ include/avutil.pxd | 9 ++ tests/test_codec_context.py | 59 ++++++++++ 11 files changed, 392 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9321d82fa..bb3c77abf 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -45,10 +45,13 @@ Features: - ``av.dump_codecs()`` now lists every codec FFmpeg knows of rather than only those with an encoder or a decoder, so data and attachment codecs appear, matching ``ffmpeg -codecs``. Its legend gains the ``..D...`` and ``..T...`` media types. - ``ContainerFormat.fixed_framesize`` reports whether a format wants fixed size audio frames. +- :class:`.CodecContext` exposes more of ``AVCodecContext``: ``pkt_timebase``, ``frame_num``, ``active_thread_type``, ``bits_per_raw_sample``, ``compression_level``, ``rc_buffer_size``, ``min_bit_rate``, a setter for ``max_bit_rate``, the audio ``initial_padding``, ``trailing_padding``, and ``seek_preroll``, and ``stats_in``/``stats_out`` for two-pass encoding. ``VideoCodecContext`` gains ``chroma_sample_location``, ``refs``, and ``mb_decision``; ``AudioCodecContext`` gains ``block_align``. +- ``CodecContext.coded_side_data`` and ``CodecContext.decoded_side_data`` expose the context's global side data as dicts of ``bytes``, keyed by packet side data name and :class:`~av.sidedata.sidedata.Type` respectively. Stream wide HDR metadata, such as mastering display and content light level, arrives in ``decoded_side_data`` once a frame has been decoded. - Enums gained the members FFmpeg has since added: ``Properties.FIELDS``, ``Properties.ENHANCEMENT``, ``PixFmtLoss.EXCESS_RESOLUTION``, ``PixFmtLoss.EXCESS_DEPTH``, ``Flags2.icc_profiles``, ``format.Flags.experimental``, ``Interpolation.STRICT``, ``Interpolation.UNSTABLE``, ``ColorTrc.V_LOG``, ``ColorPrimaries.V_GAMUT``, and the ``LCEVC``, ``VIEW_ID``, ``THREE_D_REFERENCE_DISPLAYS``, and ``EXIF`` members of ``sidedata.Type``. Fixes: +- ``CodecContext.bit_rate_tolerance`` returns its value instead of always ``None``; the getter was missing its ``return``. - A rejected ``add_stream()`` or ``add_mux_stream()`` no longer breaks the container. - ``InputContainer.size`` returns ``None`` when the size cannot be determined rather than the negative ``AVERROR`` it was passing through, which read as a plausible byte count. A non-seekable input, such as a pipe, reported ``-78``. - ``av.dump_codecs()`` no longer drops the canonical names ``h264``, ``hevc``, ``av1``, ``dirac``, and ``ilbc``, each of which was overwritten by the row of whichever encoder it resolved to. diff --git a/av/audio/codeccontext.py b/av/audio/codeccontext.py index 97180f744..08206a63c 100644 --- a/av/audio/codeccontext.py +++ b/av/audio/codeccontext.py @@ -109,3 +109,14 @@ def format(self, value): self._assert_not_open("format") format: AudioFormat = AudioFormat(value) self.ptr.sample_fmt = format.sample_fmt + + @property + def block_align(self): + """ + Number of bytes per coded audio frame, for formats with a fixed one. + + Wraps :ffmpeg:`AVCodecContext.block_align`. + + :type: int + """ + return self.ptr.block_align diff --git a/av/audio/codeccontext.pyi b/av/audio/codeccontext.pyi index e871db698..80c93afc8 100644 --- a/av/audio/codeccontext.pyi +++ b/av/audio/codeccontext.pyi @@ -25,6 +25,8 @@ class AudioCodecContext(CodecContext): layout: _Layout @property def channels(self) -> int: ... + @property + def block_align(self) -> int: ... def encode(self, frame: AudioFrame | None = None) -> list[Packet]: ... def encode_lazy(self, frame: AudioFrame | None = None) -> Iterator[Packet]: ... def decode(self, packet: Packet | None = None) -> list[AudioFrame]: ... diff --git a/av/codec/context.pxd b/av/codec/context.pxd index 2125c980e..da5e79923 100644 --- a/av/codec/context.pxd +++ b/av/codec/context.pxd @@ -16,6 +16,7 @@ cdef class CodecContext: cdef public dict options cdef HWAccel hwaccel_ctx cdef Frame _next_frame + cdef bytes _stats_in # keeps the buffer ptr.stats_in points at alive cdef uint8_t _ctxflags # ctxEnum: template_initialized # True when created via add_stream_from_template(); start_encoding() skips diff --git a/av/codec/context.py b/av/codec/context.py index bef98e20d..86803a961 100644 --- a/av/codec/context.py +++ b/av/codec/context.py @@ -10,15 +10,25 @@ from cython.cimports.av.packet import Packet from cython.cimports.av.rational import from_avrational from cython.cimports.av.utils import to_avrational +from cython.cimports.cpython.bytes import ( + PyBytes_FromString, + PyBytes_FromStringAndSize, +) from cython.cimports.libc.errno import EAGAIN from cython.cimports.libc.stdint import uint8_t from cython.cimports.libc.string import memcpy, strcmp from av.error import InvalidDataError +from av.packet import packet_sidedata_type_to_literal _cinit_sentinel = cython.declare(object, object()) +@cython.cfunc +def _to_bytes(data: cython.pointer[uint8_t], size: cython.size_t) -> bytes: + return PyBytes_FromStringAndSize(cython.cast(cython.p_char, data), size) + + @cython.cfunc def wrap_codec_context( c_ctx: cython.pointer[lib.AVCodecContext], @@ -908,14 +918,61 @@ def bit_rate(self, value: cython.longlong): @property def max_bit_rate(self): + """Maximum bitrate, or ``None`` if unset. + + Wraps :ffmpeg:`AVCodecContext.rc_max_rate`. + """ if self.ptr.rc_max_rate > 0: return self.ptr.rc_max_rate else: return None + @max_bit_rate.setter + def max_bit_rate(self, value: cython.longlong): + self.ptr.rc_max_rate = value + + @property + def min_bit_rate(self): + """Minimum bitrate, or ``None`` if unset. + + Wraps :ffmpeg:`AVCodecContext.rc_min_rate`. + """ + if self.ptr.rc_min_rate > 0: + return self.ptr.rc_min_rate + else: + return None + + @min_bit_rate.setter + def min_bit_rate(self, value: cython.longlong): + self.ptr.rc_min_rate = value + + @property + def rc_buffer_size(self): + """Decoder bitstream buffer size (VBV), in bits. + + Wraps :ffmpeg:`AVCodecContext.rc_buffer_size`. + """ + return self.ptr.rc_buffer_size + + @rc_buffer_size.setter + def rc_buffer_size(self, value: cython.int): + self.ptr.rc_buffer_size = value + + @property + def compression_level(self): + """Codec-defined compression level; ``-1`` means default. + + Wraps :ffmpeg:`AVCodecContext.compression_level`. + """ + return self.ptr.compression_level + + @compression_level.setter + def compression_level(self, value: cython.int): + self.ptr.compression_level = value + @property def bit_rate_tolerance(self): - self.ptr.bit_rate_tolerance + return self.ptr.bit_rate_tolerance @bit_rate_tolerance.setter def bit_rate_tolerance(self, value: cython.int): @@ -956,6 +1013,15 @@ def thread_type(self, value): else: self.ptr.thread_type = value.value + @property + def active_thread_type(self): + """The threading actually in use, which may differ from + :attr:`thread_type` once the codec is open. + + Wraps :ffmpeg:`AVCodecContext.active_thread_type`. + """ + return ThreadType(self.ptr.active_thread_type) + @property def skip_frame(self): """Returns one of the following str literals: @@ -1014,3 +1080,147 @@ def delay(self): """ return self.ptr.delay + + @property + def pkt_timebase(self): + """Timebase of the packets fed to this context, as a + :class:`~fractions.Fraction`. + + Decoders use it to set :attr:`.Frame.time_base`. Containers set it for + you; set it yourself when driving a bare CodecContext. + + Wraps :ffmpeg:`AVCodecContext.pkt_timebase`. + """ + return from_avrational(self.ptr.pkt_timebase) + + @pkt_timebase.setter + def pkt_timebase(self, value): + to_avrational(value, cython.address(self.ptr.pkt_timebase)) + + @property + def frame_num(self): + """Number of frames passed to/from this context so far. + + Wraps :ffmpeg:`AVCodecContext.frame_num`. + """ + return self.ptr.frame_num + + @property + def bits_per_raw_sample(self): + """Bit depth of the samples/components before encoding, e.g. ``10`` for + 10-bit video. ``0`` when unknown. + + This is the real bit depth; :attr:`.VideoCodecContext.bits_per_coded_sample` + is how many bits the bitstream spends on it. + + Wraps :ffmpeg:`AVCodecContext.bits_per_raw_sample`. + """ + return self.ptr.bits_per_raw_sample + + @bits_per_raw_sample.setter + def bits_per_raw_sample(self, value: cython.int): + self.ptr.bits_per_raw_sample = value + + @property + def initial_padding(self): + """Audio only. Samples the decoder should skip at the start of the + stream, i.e. the encoder delay. Needed for gapless playback. + + Wraps :ffmpeg:`AVCodecContext.initial_padding`. + """ + return self.ptr.initial_padding + + @property + def trailing_padding(self): + """Audio only. Samples to discard at the end of the stream. + + Wraps :ffmpeg:`AVCodecContext.trailing_padding`. + """ + return self.ptr.trailing_padding + + @trailing_padding.setter + def trailing_padding(self, value: cython.int): + self.ptr.trailing_padding = value + + @property + def seek_preroll(self): + """Number of samples to decode before the target seek point for the + output to be correct, in ``1 / AV_TIME_BASE`` units. + + Wraps :ffmpeg:`AVCodecContext.seek_preroll`. + """ + return self.ptr.seek_preroll + + @property + def stats_out(self): + """Pass-one statistics produced by the encoder, or ``None``. + + Concatenate this after every :meth:`encode` call of the first pass and + feed the result back as :attr:`stats_in` on the second. + + Wraps :ffmpeg:`AVCodecContext.stats_out`. + """ + if self.ptr.stats_out == cython.NULL: + return None + return PyBytes_FromString(self.ptr.stats_out).decode("utf-8", "replace") + + @property + def stats_in(self): + """Pass-one statistics to feed the second pass of a two-pass encode. + + Must be set before :meth:`open`. + + Wraps :ffmpeg:`AVCodecContext.stats_in`. + """ + if self.ptr.stats_in == cython.NULL: + return None + return PyBytes_FromString(self.ptr.stats_in).decode("utf-8", "replace") + + @stats_in.setter + def stats_in(self, value): + self._assert_not_open("stats_in") + if value is None: + self._stats_in = None + self.ptr.stats_in = cython.NULL + return + + # libavcodec never frees stats_in, so we keep the bytes alive ourselves. + self._stats_in = value.encode("utf-8") if type(value) is str else bytes(value) + self.ptr.stats_in = self._stats_in + + @property + def coded_side_data(self): + """Global side data attached to the coded bitstream, as a + ``dict`` of packet side data name to ``bytes``. + + Wraps :ffmpeg:`AVCodecContext.coded_side_data`. + """ + i: cython.int + return { + packet_sidedata_type_to_literal( + self.ptr.coded_side_data[i].type + ): _to_bytes( + self.ptr.coded_side_data[i].data, self.ptr.coded_side_data[i].size + ) + for i in range(self.ptr.nb_coded_side_data) + } + + @property + def decoded_side_data(self): + """Global side data produced by the decoder, as a ``dict`` of + :class:`av.sidedata.sidedata.Type` to ``bytes``. + + This is where stream-wide HDR metadata (mastering display, content + light level) shows up after the first frame is decoded. + + Wraps :ffmpeg:`AVCodecContext.decoded_side_data`. + """ + from av.sidedata.sidedata import Type + + i: cython.int + return { + Type(self.ptr.decoded_side_data[i].type): _to_bytes( + self.ptr.decoded_side_data[i].data, self.ptr.decoded_side_data[i].size + ) + for i in range(self.ptr.nb_decoded_side_data) + } diff --git a/av/codec/context.pyi b/av/codec/context.pyi index c4062f26a..d34d5b2d0 100644 --- a/av/codec/context.pyi +++ b/av/codec/context.pyi @@ -5,8 +5,9 @@ from typing import ClassVar, Literal, cast, overload from av.audio import _AudioCodecName from av.audio.codeccontext import AudioCodecContext -from av.packet import Packet +from av.packet import Packet, PktSideDataT from av.rational import AVRational +from av.sidedata.sidedata import Type as FrameSideDataType from av.subtitles import _SubtitleCodecName from av.subtitles.codeccontext import SubtitleCodecContext from av.video import _VideoCodecName @@ -134,6 +135,11 @@ class CodecContext: global_quality: int bit_rate: int | None bit_rate_tolerance: int + rc_buffer_size: int + compression_level: int + bits_per_raw_sample: int + trailing_padding: int + stats_in: str | None thread_count: int thread_type: ThreadType skip_frame: Literal[ @@ -153,6 +159,30 @@ class CodecContext: def codec(self) -> Codec: ... @property def max_bit_rate(self) -> int | None: ... + @max_bit_rate.setter + def max_bit_rate(self, value: int) -> None: ... + @property + def min_bit_rate(self) -> int | None: ... + @min_bit_rate.setter + def min_bit_rate(self, value: int) -> None: ... + @property + def pkt_timebase(self) -> AVRational: ... + @pkt_timebase.setter + def pkt_timebase(self, value: AVRational | Fraction | int) -> None: ... + @property + def frame_num(self) -> int: ... + @property + def active_thread_type(self) -> ThreadType: ... + @property + def initial_padding(self) -> int: ... + @property + def seek_preroll(self) -> int: ... + @property + def stats_out(self) -> str | None: ... + @property + def coded_side_data(self) -> dict[PktSideDataT, bytes]: ... + @property + def decoded_side_data(self) -> dict[FrameSideDataType, bytes]: ... @property def delay(self) -> bool: ... @property diff --git a/av/video/codeccontext.py b/av/video/codeccontext.py index 4c003e869..065f1d993 100644 --- a/av/video/codeccontext.py +++ b/av/video/codeccontext.py @@ -496,3 +496,49 @@ def qmax(self): @qmax.setter def qmax(self, value): self.ptr.qmax = value + + @property + def chroma_sample_location(self): + """ + Location of the chroma samples relative to the luma samples, as + FFmpeg's raw integer value. + + Wraps :ffmpeg:`AVCodecContext.chroma_sample_location`. + + :type: int + """ + return self.ptr.chroma_sample_location + + @chroma_sample_location.setter + def chroma_sample_location(self, value: cython.int): + self.ptr.chroma_sample_location = cython.cast(lib.AVChromaLocation, value) + + @property + def refs(self): + """ + The number of reference frames. + + Wraps :ffmpeg:`AVCodecContext.refs`. + + :type: int + """ + return self.ptr.refs + + @refs.setter + def refs(self, value: cython.int): + self.ptr.refs = value + + @property + def mb_decision(self): + """ + Macroblock decision mode: 0 (simple), 1 (bits) or 2 (rate distortion). + + Wraps :ffmpeg:`AVCodecContext.mb_decision`. + + :type: int + """ + return self.ptr.mb_decision + + @mb_decision.setter + def mb_decision(self, value: cython.int): + self.ptr.mb_decision = value diff --git a/av/video/codeccontext.pyi b/av/video/codeccontext.pyi index 056ae4629..27d902c75 100644 --- a/av/video/codeccontext.pyi +++ b/av/video/codeccontext.pyi @@ -43,6 +43,9 @@ class VideoCodecContext(CodecContext): color_trc: int colorspace: int field_order: int + chroma_sample_location: int + refs: int + mb_decision: int qmin: int qmax: int type: Literal["video"] diff --git a/include/avcodec.pxd b/include/avcodec.pxd index ae1afafc9..5d08c6525 100644 --- a/include/avcodec.pxd +++ b/include/avcodec.pxd @@ -248,6 +248,7 @@ cdef extern from "libavcodec/avcodec.h" nogil: AVColorTransferCharacteristic color_trc AVColorSpace colorspace AVColorRange color_range + AVChromaLocation chroma_sample_location AVFieldOrder field_order int has_b_frames @@ -259,6 +260,10 @@ cdef extern from "libavcodec/avcodec.h" nogil: AVSampleFormat sample_fmt AVChannelLayout ch_layout int frame_size + int block_align + int initial_padding + int trailing_padding + int seek_preroll int bit_rate_tolerance int global_quality @@ -268,6 +273,8 @@ cdef extern from "libavcodec/avcodec.h" nogil: int rc_buffer_size int64_t rc_max_rate int64_t rc_min_rate + char *stats_in + char *stats_out const AVHWAccel *hwaccel AVBufferRef *hw_device_ctx @@ -275,7 +282,10 @@ cdef extern from "libavcodec/avcodec.h" nogil: int thread_count int thread_type + int active_thread_type int bits_per_coded_sample + int bits_per_raw_sample + int refs int profile int level AVDiscard skip_frame @@ -284,6 +294,11 @@ cdef extern from "libavcodec/avcodec.h" nogil: uint8_t *subtitle_header int64_t frame_num + AVPacketSideData *coded_side_data + int nb_coded_side_data + AVFrameSideData **decoded_side_data + int nb_decoded_side_data + cdef AVCodecContext* avcodec_alloc_context3(const AVCodec *codec) cdef void avcodec_free_context(AVCodecContext **ctx) cdef const AVCodec* avcodec_find_decoder(AVCodecID id) @@ -525,3 +540,4 @@ cdef extern from "libavcodec/packet.h" nogil: AVPacket *pkt, AVPacketSideDataType type, uint8_t *data, size_t size ) const char *av_packet_side_data_name(AVPacketSideDataType type) + const char *av_frame_side_data_name(AVFrameSideDataType type) diff --git a/include/avutil.pxd b/include/avutil.pxd index 034635b7c..b7aa514c5 100644 --- a/include/avutil.pxd +++ b/include/avutil.pxd @@ -62,6 +62,15 @@ cdef extern from "libavutil/avutil.h" nogil: AVCOL_PRI_EBU3213 AVCOL_PRI_V_GAMUT + cdef enum AVChromaLocation: + AVCHROMA_LOC_UNSPECIFIED + AVCHROMA_LOC_LEFT + AVCHROMA_LOC_CENTER + AVCHROMA_LOC_TOPLEFT + AVCHROMA_LOC_TOP + AVCHROMA_LOC_BOTTOMLEFT + AVCHROMA_LOC_BOTTOM + cdef enum AVColorTransferCharacteristic: AVCOL_TRC_BT709 AVCOL_TRC_UNSPECIFIED diff --git a/tests/test_codec_context.py b/tests/test_codec_context.py index 272a1d45a..44b06eab7 100644 --- a/tests/test_codec_context.py +++ b/tests/test_codec_context.py @@ -617,3 +617,62 @@ def _audio_encoding( result_samples += frame.samples assert frame.sample_rate == sample_rate assert frame.layout.nb_channels == 2 + + +class TestNewlyExposedFields(TestCase): + def test_encoder_scalars_roundtrip(self) -> None: + ctx = av.CodecContext.create("libx264", "w") + assert isinstance(ctx, av.video.codeccontext.VideoCodecContext) + + ctx.bit_rate_tolerance = 4000 + ctx.max_bit_rate = 2_000_000 + ctx.min_bit_rate = 500_000 + ctx.rc_buffer_size = 1_000_000 + ctx.compression_level = 5 + ctx.bits_per_raw_sample = 8 + ctx.trailing_padding = 0 + ctx.chroma_sample_location = 1 + ctx.refs = 3 + ctx.mb_decision = 2 + ctx.pkt_timebase = Fraction(1, 1000) + + assert ctx.bit_rate_tolerance == 4000 + assert ctx.max_bit_rate == 2_000_000 + assert ctx.min_bit_rate == 500_000 + assert ctx.rc_buffer_size == 1_000_000 + assert ctx.compression_level == 5 + assert ctx.bits_per_raw_sample == 8 + assert ctx.trailing_padding == 0 + assert ctx.chroma_sample_location == 1 + assert ctx.refs == 3 + assert ctx.mb_decision == 2 + assert ctx.pkt_timebase == Fraction(1, 1000) + assert ctx.frame_num == 0 + + def test_stats_in_out(self) -> None: + ctx = av.CodecContext.create("libx264", "w") + assert ctx.stats_in is None + assert ctx.stats_out is None + + ctx.stats_in = "frame in:0 out:0;" + assert ctx.stats_in == "frame in:0 out:0;" + ctx.stats_in = None + assert ctx.stats_in is None + + def test_audio_padding(self) -> None: + with av.open(fate_suite("mkv/codec_delay_opus.mkv")) as container: + ctx = container.streams.audio[0].codec_context + assert ctx.initial_padding == 312 + assert ctx.seek_preroll == 3840 + assert ctx.block_align == 0 + + def test_coded_side_data(self) -> None: + with av.open(fate_suite("mov/displaymatrix.mov")) as container: + stream = container.streams.video[0] + ctx = stream.codec_context + assert ctx.bits_per_raw_sample == 8 + assert len(ctx.coded_side_data["display_matrix"]) == 36 + + for _ in container.decode(stream): + break + assert ctx.frame_num == 1