Add add_stream_from_template() method

av/container/output.pyi

@@ -1,14 +1,15 @@
 from fractions import Fraction
-from typing import Literal, Sequence, overload
+from typing import Literal, Sequence, TypeVar, overload
 
-from av.audio.layout import AudioLayout
 from av.audio.stream import AudioStream
 from av.packet import Packet
 from av.stream import Stream
 from av.video.stream import VideoStream
 
 from .core import Container
 
+_StreamT = TypeVar("_StreamT", bound=Stream, default=Stream)
+
 class OutputContainer(Container):
     def __enter__(self) -> OutputContainer: ...
     @overload
@@ -30,6 +31,15 @@ class OutputContainer(Container):
         **kwargs,
     ) -> VideoStream: ...
     @overload
+    def add_stream(
+        self,
+        codec_name: None = None,
+        rate: Fraction | int | None = None,
+        template: _StreamT | None = None,
+        options: dict[str, str] | None = None,
+        **kwargs,
+    ) -> _StreamT: ...
+    @overload
     def add_stream(
         self,
         codec_name: str | None = None,
@@ -38,6 +48,7 @@ class OutputContainer(Container):
         options: dict[str, str] | None = None,
         **kwargs,
     ) -> Stream: ...
+    def add_stream_from_template(self, template: _StreamT, **kwargs) -> _StreamT: ...
     def start_encoding(self) -> None: ...
     def close(self) -> None: ...
     def mux(self, packets: Packet | Sequence[Packet]) -> None: ...

av/container/output.pyx

@@ -46,16 +46,15 @@ cdef class OutputContainer(Container):
     def add_stream(self, codec_name=None, object rate=None, Stream template=None, options=None, **kwargs):
         """add_stream(codec_name, rate=None)
 
-        Create a new stream, and return it.
+        Creates a new stream from a codec name or a template, and returns it.
 
-        :param str codec_name: The name of a codec.
-        :param rate: The frame rate for video, and sample rate for audio.
-            Examples for video include ``24`` and ``Fraction(30000, 1001)``.
-            Examples for audio include ``48000`` and ``44100``.
+        :param codec_name: The name of a codec.
+        :type codec_name: str | Codec | None
         :param template: Copy codec from another :class:`~av.stream.Stream` instance.
+        :type template: :class:`~av.stream.Stream` | None
         :param dict options: Stream options.
-        :param \\**kwargs: Set attributes of the stream.
-        :returns: The new :class:`~av.stream.Stream`.
+        :param \\**kwargs: Set attributes for the stream.
+        :rtype: The new :class:`~av.stream.Stream`.
 
         """
 
@@ -76,7 +75,7 @@ cdef class OutputContainer(Container):
         # 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"
+                f"{self.format.name!r} format does not support {codec_obj.name!r} codec"
             )
 
         # Create new stream in the AVFormatContext, set AVCodecContext values.
@@ -147,6 +146,63 @@ cdef class OutputContainer(Container):
 
         return py_stream
 
+    def add_stream_from_template(self, Stream template not None, **kwargs):
+        """
+        Creates a new stream from a template.
+
+        :param template: Copy codec from another :class:`~av.stream.Stream` instance.
+        :param \\**kwargs: Set attributes for the stream.
+        :rtype: The new :class:`~av.stream.Stream`.
+
+        """
+
+        if not template.codec_context:
+            raise ValueError("template has no codec context")
+
+        cdef Codec codec_obj = Codec(template.codec_context.codec.name, "w")
+        cdef const lib.AVCodec *codec = 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):
+            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)
+
+        err_check(lib.avcodec_parameters_to_context(codec_context, template.ptr.codecpar))
+        # Reset the codec tag assuming we are remuxing.
+        codec_context.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
+
+        # 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))
+
+        # Construct the user-land stream
+        cdef CodecContext 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 template.type == "video":
+            py_stream.time_base = kwargs.pop("time_base", 1 / template.average_rate)
+        elif template.type == "audio":
+            py_stream.time_base = kwargs.pop("time_base", 1 / template.rate)
+        else:
+            py_stream.time_base = kwargs.pop("time_base", None)
+
+        for k, v in kwargs.items():
+            setattr(py_stream, k, v)
+
+        return py_stream
+
     cpdef start_encoding(self):
         """Write the file header! Called automatically."""
 

tests/test_streams.py

@@ -59,6 +59,25 @@ def test_printing_video_stream(self) -> None:
         container.close()
         input_.close()
 
+    def test_printing_video_stream2(self) -> None:
+        input_ = av.open(fate_suite("h264/interlaced_crop.mp4"))
+        input_stream = input_.streams.video[0]
+        container = av.open("out.mkv", "w")
+
+        video_stream = container.add_stream_from_template(input_stream)
+        encoder = video_stream.codec.name
+
+        for frame in input_.decode(video=0):
+            container.mux(video_stream.encode(frame))
+            break
+
+        repr = f"{video_stream}"
+        assert repr.startswith(f"<av.VideoStream #0 {encoder}, yuv420p 640x360 at ")
+        assert repr.endswith(">")
+
+        container.close()
+        input_.close()
+
     # def test_side_data(self) -> None:
     #     container = av.open(fate_suite("mov/displaymatrix.mov"))
     #     video = container.streams.video[0]