From 73db001a36299332c7a6419e4252076b289cd78f Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Sat, 16 Nov 2024 03:34:00 -0500 Subject: [PATCH] Add `add_data_stream()` method --- av/container/output.pyi | 11 +++++-- av/container/output.pyx | 67 +++++++++++++++++++++++++++++++++++++++-- tests/test_streams.py | 38 +++++++++++++++++++++++ 3 files changed, 111 insertions(+), 5 deletions(-) diff --git a/av/container/output.pyi b/av/container/output.pyi index 97b67ee71..a7c89452f 100644 --- a/av/container/output.pyi +++ b/av/container/output.pyi @@ -1,14 +1,16 @@ from fractions import Fraction -from typing import Literal, Sequence, TypeVar, overload +from typing import Literal, Sequence, TypeVar, Union, overload from av.audio.stream import AudioStream +from av.data.stream import DataStream from av.packet import Packet from av.stream import Stream +from av.subtitles.stream import SubtitleStream from av.video.stream import VideoStream from .core import Container -_StreamT = TypeVar("_StreamT", bound=Stream, default=Stream) +_StreamT = TypeVar("_StreamT", bound=Union[VideoStream, AudioStream, SubtitleStream]) class OutputContainer(Container): def __enter__(self) -> OutputContainer: ... @@ -35,8 +37,11 @@ class OutputContainer(Container): rate: Fraction | int | None = None, options: dict[str, str] | None = None, **kwargs, - ) -> Stream: ... + ) -> VideoStream | AudioStream | SubtitleStream: ... def add_stream_from_template(self, template: _StreamT, **kwargs) -> _StreamT: ... + def add_data_stream( + self, codec_name: str | None = None, options: dict[str, str] | None = None + ) -> DataStream: ... def start_encoding(self) -> None: ... def close(self) -> None: ... def mux(self, packets: Packet | Sequence[Packet]) -> None: ... diff --git a/av/container/output.pyx b/av/container/output.pyx index 3e4c9dbb4..a07bc79e8 100644 --- a/av/container/output.pyx +++ b/av/container/output.pyx @@ -47,6 +47,7 @@ cdef class OutputContainer(Container): """add_stream(codec_name, rate=None) Creates a new stream from a codec name and returns it. + Supports video, audio, and subtitle streams. :param codec_name: The name of a codec. :type codec_name: str | Codec @@ -137,7 +138,7 @@ cdef class OutputContainer(Container): def add_stream_from_template(self, Stream template not None, **kwargs): """ - Creates a new stream from a template. + Creates a new stream from a template. Supports video, audio, and subtitle streams. :param template: Copy codec from another :class:`~av.stream.Stream` instance. :param \\**kwargs: Set attributes for the stream. @@ -192,6 +193,65 @@ cdef class OutputContainer(Container): return py_stream + + def add_data_stream(self, codec_name=None, dict options=None): + """add_data_stream(codec_name=None) + + Creates a new data stream and returns it. + + :param codec_name: Optional name of the data codec (e.g. 'klv') + :type codec_name: str | None + :param dict options: Stream options. + :rtype: The new :class:`~av.data.stream.DataStream`. + """ + cdef const lib.AVCodec *codec = NULL + + if codec_name is not None: + codec = lib.avcodec_find_encoder_by_name(codec_name.encode()) + if codec == 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): + 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: + 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: + 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 + + # Initialize stream codec parameters + err_check(lib.avcodec_parameters_from_context(stream.codecpar, codec_context)) + 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) + + cdef Stream py_stream = wrap_stream(self, stream, py_codec_context) + self.streams.add_stream(py_stream) + + if options: + py_stream.options.update(options) + + return py_stream + cpdef start_encoding(self): """Write the file header! Called automatically.""" @@ -206,8 +266,11 @@ cdef class OutputContainer(Container): cdef Stream stream for stream in self.streams: ctx = stream.codec_context + # Skip codec context handling for data streams without codecs if ctx is None: - raise ValueError(f"Stream {stream.index} has no codec context") + if stream.type != "data": + raise ValueError(f"Stream {stream.index} has no codec context") + continue if not ctx.is_open: for k, v in self.options.items(): diff --git a/tests/test_streams.py b/tests/test_streams.py index 045745543..b1f2bd12f 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -1,3 +1,5 @@ +from fractions import Fraction + import av from .common import fate_suite @@ -78,6 +80,42 @@ def test_printing_video_stream2(self) -> None: container.close() input_.close() + def test_data_stream(self) -> None: + # First test writing and reading a simple data stream + container1 = av.open("data.ts", "w") + data_stream = container1.add_data_stream() + + test_data = [b"test data 1", b"test data 2", b"test data 3"] + for i, data_ in enumerate(test_data): + packet = av.Packet(data_) + packet.pts = i + packet.stream = data_stream + container1.mux(packet) + container1.close() + + # Test reading back the data stream + container = av.open("data.ts") + + # Test best stream selection + data = container.streams.best("data") + assert data == container.streams.data[0] + + # Test get method + assert [data] == container.streams.get(data=0) + assert [data] == container.streams.get(data=(0,)) + + # Verify we can read back all the packets, ignoring empty ones + packets = [p for p in container.demux(data) if bytes(p)] + assert len(packets) == len(test_data) + for packet, original_data in zip(packets, test_data): + assert bytes(packet) == original_data + + # Test string representation + repr = f"{data_stream}" + assert repr.startswith("") + + container.close() + # def test_side_data(self) -> None: # container = av.open(fate_suite("mov/displaymatrix.mov")) # video = container.streams.video[0]