From 82bc9e1ae7fd8c0125fed25915e6154576fd0c04 Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Tue, 20 May 2025 23:42:54 -0400 Subject: [PATCH] Make container/output pure This also syncs the QoL changes I've made in basswood-av. "Pure" stands for Cython's "Pure Python" mode, a way to write Cython in the Python grammar. --- av/container/output.pxd | 1 - av/container/{output.pyx => output.py} | 236 +++++++++++++------------ 2 files changed, 127 insertions(+), 110 deletions(-) rename av/container/{output.pyx => output.py} (58%) diff --git a/av/container/output.pxd b/av/container/output.pxd index a4299891c..51d3f308e 100644 --- a/av/container/output.pxd +++ b/av/container/output.pxd @@ -5,7 +5,6 @@ from av.stream cimport Stream cdef class OutputContainer(Container): - cdef bint _started cdef bint _done cdef lib.AVPacket *packet_ptr diff --git a/av/container/output.pyx b/av/container/output.py similarity index 58% rename from av/container/output.pyx rename to av/container/output.py index 2a12c40f2..64a4f4a53 100644 --- a/av/container/output.pyx +++ b/av/container/output.py @@ -1,24 +1,22 @@ -import logging import os from fractions import Fraction -cimport libav as lib - -from av.codec.codec cimport Codec -from av.codec.context cimport CodecContext, wrap_codec_context -from av.container.streams cimport StreamContainer -from av.dictionary cimport _Dictionary -from av.error cimport err_check -from av.packet cimport Packet -from av.stream cimport Stream, wrap_stream -from av.utils cimport dict_to_avdict, to_avrational +import cython +from cython.cimports import libav as lib +from cython.cimports.av.codec.codec import Codec +from cython.cimports.av.codec.context import CodecContext, wrap_codec_context +from cython.cimports.av.container.streams import StreamContainer +from cython.cimports.av.dictionary import _Dictionary +from cython.cimports.av.error import err_check +from cython.cimports.av.packet import Packet +from cython.cimports.av.stream import Stream, wrap_stream +from cython.cimports.av.utils import dict_to_avdict, to_avrational from av.dictionary import Dictionary -log = logging.getLogger(__name__) - -cdef close_output(OutputContainer self): +@cython.cfunc +def close_output(self: OutputContainer): self.streams = StreamContainer() if self._started and not self._done: # We must only ever call av_write_trailer *once*, otherwise we get a @@ -28,23 +26,24 @@ self.err_check(lib.av_write_trailer(self.ptr)) finally: if self.file is None and not (self.ptr.oformat.flags & lib.AVFMT_NOFILE): - lib.avio_closep(&self.ptr.pb) + lib.avio_closep(cython.address(self.ptr.pb)) self._done = True -cdef class OutputContainer(Container): +@cython.cclass +class OutputContainer(Container): def __cinit__(self, *args, **kwargs): self.streams = StreamContainer() self.metadata = {} - with nogil: + with cython.nogil: self.packet_ptr = lib.av_packet_alloc() def __dealloc__(self): close_output(self) - with nogil: - lib.av_packet_free(&self.packet_ptr) + with cython.nogil: + lib.av_packet_free(cython.address(self.packet_ptr)) - def add_stream(self, codec_name, rate=None, dict options=None, **kwargs): + def add_stream(self, codec_name, rate=None, options: dict | None = None, **kwargs): """add_stream(codec_name, rate=None) Creates a new stream from a codec name and returns it. @@ -58,67 +57,69 @@ def add_stream(self, codec_name, rate=None, dict options=None, **kwargs): """ - cdef Codec codec_obj = Codec(codec_name, "w") - cdef const lib.AVCodec *codec = codec_obj.ptr + codec_obj: Codec = Codec(codec_name, "w") + codec: cython.pointer[cython.const[lib.AVCodec]] = codec_obj.ptr # Assert that this format supports the requested codec. - if not lib.avformat_query_codec(self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL): + if not lib.avformat_query_codec( + self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL + ): raise ValueError( f"{self.format.name!r} format does not support {codec_obj.name!r} codec" ) # Create new stream in the AVFormatContext, set AVCodecContext values. - cdef lib.AVStream *stream = lib.avformat_new_stream(self.ptr, codec) - cdef lib.AVCodecContext *codec_context = lib.avcodec_alloc_context3(codec) + stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(self.ptr, codec) + ctx: cython.pointer[lib.AVCodecContext] = lib.avcodec_alloc_context3(codec) # Now lets set some more sane video defaults if codec.type == lib.AVMEDIA_TYPE_VIDEO: - codec_context.pix_fmt = lib.AV_PIX_FMT_YUV420P - codec_context.width = kwargs.pop("width", 640) - codec_context.height = kwargs.pop("height", 480) - codec_context.bit_rate = kwargs.pop("bit_rate", 0) - codec_context.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 128000) + ctx.pix_fmt = lib.AV_PIX_FMT_YUV420P + ctx.width = kwargs.pop("width", 640) + ctx.height = kwargs.pop("height", 480) + ctx.bit_rate = kwargs.pop("bit_rate", 0) + ctx.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 128000) try: - to_avrational(kwargs.pop("time_base"), &codec_context.time_base) + to_avrational(kwargs.pop("time_base"), cython.address(ctx.time_base)) except KeyError: pass - to_avrational(rate or 24, &codec_context.framerate) + to_avrational(rate or 24, cython.address(ctx.framerate)) - stream.avg_frame_rate = codec_context.framerate - stream.time_base = codec_context.time_base + stream.avg_frame_rate = ctx.framerate + stream.time_base = ctx.time_base # Some sane audio defaults elif codec.type == lib.AVMEDIA_TYPE_AUDIO: - codec_context.sample_fmt = codec.sample_fmts[0] - codec_context.bit_rate = kwargs.pop("bit_rate", 0) - codec_context.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 32000) + ctx.sample_fmt = codec.sample_fmts[0] + ctx.bit_rate = kwargs.pop("bit_rate", 0) + ctx.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 32000) try: - to_avrational(kwargs.pop("time_base"), &codec_context.time_base) + to_avrational(kwargs.pop("time_base"), cython.address(ctx.time_base)) except KeyError: pass if rate is None: - codec_context.sample_rate = 48000 + ctx.sample_rate = 48000 elif type(rate) is int: - codec_context.sample_rate = rate + ctx.sample_rate = rate else: raise TypeError("audio stream `rate` must be: int | None") - stream.time_base = codec_context.time_base - lib.av_channel_layout_default(&codec_context.ch_layout, 2) + stream.time_base = ctx.time_base + lib.av_channel_layout_default(cython.address(ctx.ch_layout), 2) # Some formats want stream headers to be separate if self.ptr.oformat.flags & lib.AVFMT_GLOBALHEADER: - codec_context.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER + ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER # Initialise stream codec parameters to populate the codec type. # # Subsequent changes to the codec context will be applied just before # encoding starts in `start_encoding()`. - err_check(lib.avcodec_parameters_from_context(stream.codecpar, codec_context)) + err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx)) # Construct the user-land stream - cdef CodecContext py_codec_context = wrap_codec_context(codec_context, codec, None) - cdef Stream py_stream = wrap_stream(self, stream, py_codec_context) + py_codec_context: CodecContext = wrap_codec_context(ctx, codec, None) + py_stream: Stream = wrap_stream(self, stream, py_codec_context) self.streams.add_stream(py_stream) if options: @@ -129,7 +130,7 @@ def add_stream(self, codec_name, rate=None, dict options=None, **kwargs): return py_stream - def add_stream_from_template(self, Stream template not None, opaque=None, **kwargs): + def add_stream_from_template(self, template: Stream, opaque=None, **kwargs): """ Creates a new stream from a template. Supports video, audio, and subtitle streams. @@ -138,43 +139,44 @@ def add_stream_from_template(self, Stream template not None, opaque=None, **kwar :param \\**kwargs: Set attributes for the stream. :rtype: The new :class:`~av.stream.Stream`. """ - cdef const lib.AVCodec *codec - cdef Codec codec_obj - if opaque is None: opaque = template.type != "video" + codec_obj: Codec if opaque: # Copy ctx from template. codec_obj = template.codec_context.codec - else: # Construct new codec object. + else: # Construct new codec object. codec_obj = Codec(template.codec_context.codec.name, "w") - codec = codec_obj.ptr + + codec: cython.pointer[cython.const[lib.AVCodec]] = codec_obj.ptr # Assert that this format supports the requested codec. - if not lib.avformat_query_codec(self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL): + if not lib.avformat_query_codec( + self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL + ): raise ValueError( f"{self.format.name!r} format does not support {codec_obj.name!r} codec" ) # Create new stream in the AVFormatContext, set AVCodecContext values. - cdef lib.AVStream *stream = lib.avformat_new_stream(self.ptr, codec) - cdef lib.AVCodecContext *codec_context = lib.avcodec_alloc_context3(codec) + stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(self.ptr, codec) + ctx: cython.pointer[lib.AVCodecContext] = lib.avcodec_alloc_context3(codec) - err_check(lib.avcodec_parameters_to_context(codec_context, template.ptr.codecpar)) + err_check(lib.avcodec_parameters_to_context(ctx, template.ptr.codecpar)) # Reset the codec tag assuming we are remuxing. - codec_context.codec_tag = 0 + ctx.codec_tag = 0 # Some formats want stream headers to be separate if self.ptr.oformat.flags & lib.AVFMT_GLOBALHEADER: - codec_context.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER + ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER # Initialize stream codec parameters to populate the codec type. Subsequent changes to # the codec context will be applied just before encoding starts in `start_encoding()`. - err_check(lib.avcodec_parameters_from_context(stream.codecpar, codec_context)) + err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx)) # Construct the user-land stream - cdef CodecContext py_codec_context = wrap_codec_context(codec_context, codec, None) - cdef Stream py_stream = wrap_stream(self, stream, py_codec_context) + py_codec_context: CodecContext = wrap_codec_context(ctx, codec, None) + py_stream: Stream = wrap_stream(self, stream, py_codec_context) self.streams.add_stream(py_stream) for k, v in kwargs.items(): @@ -182,8 +184,7 @@ def add_stream_from_template(self, Stream template not None, opaque=None, **kwar return py_stream - - def add_data_stream(self, codec_name=None, dict options=None): + def add_data_stream(self, codec_name=None, options: dict | None = None): """add_data_stream(codec_name=None) Creates a new data stream and returns it. @@ -193,47 +194,49 @@ def add_data_stream(self, codec_name=None, dict options=None): :param dict options: Stream options. :rtype: The new :class:`~av.data.stream.DataStream`. """ - cdef const lib.AVCodec *codec = NULL + codec: cython.pointer[cython.const[lib.AVCodec]] = cython.NULL if codec_name is not None: codec = lib.avcodec_find_encoder_by_name(codec_name.encode()) - if codec == NULL: + if codec == cython.NULL: raise ValueError(f"Unknown data codec: {codec_name}") # Assert that this format supports the requested codec - if not lib.avformat_query_codec(self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL): + if not lib.avformat_query_codec( + self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL + ): raise ValueError( f"{self.format.name!r} format does not support {codec_name!r} codec" ) # Create new stream in the AVFormatContext - cdef lib.AVStream *stream = lib.avformat_new_stream(self.ptr, codec) - if stream == NULL: + stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(self.ptr, codec) + if stream == cython.NULL: raise MemoryError("Could not allocate stream") # Set up codec context if we have a codec - cdef lib.AVCodecContext *codec_context = NULL - if codec != NULL: - codec_context = lib.avcodec_alloc_context3(codec) - if codec_context == NULL: + ctx: cython.pointer[lib.AVCodecContext] = cython.NULL + if codec != cython.NULL: + ctx = lib.avcodec_alloc_context3(codec) + if ctx == cython.NULL: raise MemoryError("Could not allocate codec context") # Some formats want stream headers to be separate if self.ptr.oformat.flags & lib.AVFMT_GLOBALHEADER: - codec_context.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER + ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER # Initialize stream codec parameters - err_check(lib.avcodec_parameters_from_context(stream.codecpar, codec_context)) + err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx)) else: # For raw data streams, just set the codec type stream.codecpar.codec_type = lib.AVMEDIA_TYPE_DATA # Construct the user-land stream - cdef CodecContext py_codec_context = None - if codec_context != NULL: - py_codec_context = wrap_codec_context(codec_context, codec, None) + py_codec_context: CodecContext | None = None + if ctx != cython.NULL: + py_codec_context = wrap_codec_context(ctx, codec, None) - cdef Stream py_stream = wrap_stream(self, stream, py_codec_context) + py_stream: Stream = wrap_stream(self, stream, py_codec_context) self.streams.add_stream(py_stream) if options: @@ -241,18 +244,18 @@ def add_data_stream(self, codec_name=None, dict options=None): return py_stream - cpdef start_encoding(self): + @cython.ccall + def start_encoding(self): """Write the file header! Called automatically.""" - if self._started: return # TODO: This does NOT handle options coming from 3 sources. # This is only a rough approximation of what would be cool to do. - used_options = set() + used_options: set = set() + stream: Stream # Finalize and open all streams. - cdef Stream stream for stream in self.streams: ctx = stream.codec_context # Skip codec context handling for data streams without codecs @@ -274,29 +277,38 @@ def add_data_stream(self, codec_name=None, dict options=None): stream._finalize_for_output() # Open the output file, if needed. - cdef bytes name_obj = os.fsencode(self.name if self.file is None else "") - cdef char *name = name_obj - if self.ptr.pb == NULL and not self.ptr.oformat.flags & lib.AVFMT_NOFILE: - err_check(lib.avio_open(&self.ptr.pb, name, lib.AVIO_FLAG_WRITE)) + name_obj: bytes = os.fsencode(self.name if self.file is None else "") + name: cython.p_char = name_obj + if self.ptr.pb == cython.NULL and not self.ptr.oformat.flags & lib.AVFMT_NOFILE: + err_check( + lib.avio_open(cython.address(self.ptr.pb), name, lib.AVIO_FLAG_WRITE) + ) # Copy the metadata dict. dict_to_avdict( - &self.ptr.metadata, self.metadata, + cython.address(self.ptr.metadata), + self.metadata, encoding=self.metadata_encoding, - errors=self.metadata_errors + errors=self.metadata_errors, ) - cdef _Dictionary all_options = Dictionary(self.options, self.container_options) - cdef _Dictionary options = all_options.copy() - self.err_check(lib.avformat_write_header(self.ptr, &options.ptr)) + all_options: _Dictionary = Dictionary(self.options, self.container_options) + options: _Dictionary = all_options.copy() + self.err_check(lib.avformat_write_header(self.ptr, cython.address(options.ptr))) # Track option usage... for k in all_options: if k not in options: used_options.add(k) + # ... and warn if any weren't used. - unused_options = {k: v for k, v in self.options.items() if k not in used_options} + unused_options = { + k: v for k, v in self.options.items() if k not in used_options + } if unused_options: + import logging + + log = logging.getLogger(__name__) log.warning("Some options were not used: %s" % unused_options) self._started = True @@ -306,21 +318,25 @@ def supported_codecs(self): """ Returns a set of all codecs this format supports. """ - result = set() - cdef const lib.AVCodec *codec = NULL - cdef void *opaque = NULL + result: set = set() + codec: cython.pointer[cython.const[lib.AVCodec]] = cython.NULL + opaque: cython.p_void = cython.NULL while True: - codec = lib.av_codec_iterate(&opaque) - if codec == NULL: + codec = lib.av_codec_iterate(cython.address(opaque)) + if codec == cython.NULL: break - if lib.avformat_query_codec(self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL) == 1: + if ( + lib.avformat_query_codec( + self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL + ) + == 1 + ): result.add(codec.name) return result - @property def default_video_codec(self): """ @@ -346,29 +362,31 @@ def close(self): close_output(self) def mux(self, packets): - # We accept either a Packet, or a sequence of packets. This should - # smooth out the transition to the new encode API which returns a - # sequence of packets. + # We accept either a Packet, or a sequence of packets. This should smooth out + # the transition to the new encode API which returns a sequence of packets. if isinstance(packets, Packet): self.mux_one(packets) else: for packet in packets: self.mux_one(packet) - def mux_one(self, Packet packet not None): + def mux_one(self, packet: Packet): self.start_encoding() # Assert the packet is in stream time. - if packet.ptr.stream_index < 0 or packet.ptr.stream_index >= self.ptr.nb_streams: + if ( + packet.ptr.stream_index < 0 + or cython.cast(cython.uint, packet.ptr.stream_index) >= self.ptr.nb_streams + ): raise ValueError("Bad Packet stream_index.") - cdef lib.AVStream *stream = self.ptr.streams[packet.ptr.stream_index] + + stream: cython.pointer[lib.AVStream] = self.ptr.streams[packet.ptr.stream_index] packet._rebase_time(stream.time_base) - # Make another reference to the packet, as av_interleaved_write_frame + # Make another reference to the packet, as `av_interleaved_write_frame()` # takes ownership of the reference. self.err_check(lib.av_packet_ref(self.packet_ptr, packet.ptr)) - cdef int ret - with nogil: - ret = lib.av_interleaved_write_frame(self.ptr, self.packet_ptr) + with cython.nogil: + ret: cython.int = lib.av_interleaved_write_frame(self.ptr, self.packet_ptr) self.err_check(ret)