Writable and copyable attachment and data streams

Author: skeskinen

av/container/output.py

@@ -134,7 +134,7 @@ def add_stream_from_template(
         self, template: Stream, opaque: bool | None = None, **kwargs
     ):
         """
-        Creates a new stream from a template. Supports video, audio, and subtitle streams.
+        Creates a new stream from a template. Supports video, audio, subtitle, data and attachment streams.
 
         :param template: Copy codec from another :class:`~av.stream.Stream` instance.
         :param opaque: If True, copy opaque data from the template's codec context.
@@ -145,9 +145,7 @@ def add_stream_from_template(
             opaque = template.type != "video"
 
         if template.codec_context is None:
-            raise ValueError(
-                f"template stream of type {template.type} has no codec context"
-            )
+            return self._add_stream_without_codec_from_template(template, **kwargs)
 
         codec_obj: Codec
         if opaque:  # Copy ctx from template.
@@ -196,6 +194,79 @@ def add_stream_from_template(
 
         return py_stream
 
+    def _add_stream_without_codec_from_template(
+        self, template: Stream, **kwargs
+    ) -> Stream:
+        codec_type: cython.int = template.ptr.codecpar.codec_type
+        if codec_type not in {lib.AVMEDIA_TYPE_ATTACHMENT, lib.AVMEDIA_TYPE_DATA}:
+            raise ValueError(
+                f"template stream of type {template.type} has no codec context"
+            )
+
+        stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(
+            self.ptr, cython.NULL
+        )
+        if stream == cython.NULL:
+            raise MemoryError("Could not allocate stream")
+
+        err_check(lib.avcodec_parameters_copy(stream.codecpar, template.ptr.codecpar))
+
+        # Mirror basic properties that are not derived from a codec context.
+        stream.time_base = template.ptr.time_base
+        stream.start_time = template.ptr.start_time
+        stream.duration = template.ptr.duration
+        stream.disposition = template.ptr.disposition
+
+        py_stream: Stream = wrap_stream(self, stream, None)
+        self.streams.add_stream(py_stream)
+
+        py_stream.metadata = dict(template.metadata)
+
+        for k, v in kwargs.items():
+            setattr(py_stream, k, v)
+
+        return py_stream
+
+    def add_attachment(self, name: str, mimetype: str, data: bytes):
+        """
+        Create an attachment stream and embed its payload into the container header.
+
+        - Only supported by formats that support attachments (e.g. Matroska).
+        - No per-packet muxing is required; attachments are written at header time.
+        """
+        # Create stream with no codec (attachments are codec-less).
+        stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(
+            self.ptr, cython.NULL
+        )
+        if stream == cython.NULL:
+            raise MemoryError("Could not allocate stream")
+
+        stream.codecpar.codec_type = lib.AVMEDIA_TYPE_ATTACHMENT
+        stream.codecpar.codec_id = lib.AV_CODEC_ID_NONE
+
+        # Allocate and copy payload into codecpar.extradata.
+        payload_size: cython.size_t = len(data)
+        if payload_size:
+            buf = cython.cast(cython.p_uchar, lib.av_malloc(payload_size + 1))
+            if buf == cython.NULL:
+                raise MemoryError("Could not allocate attachment data")
+            # Copy bytes.
+            for i in range(payload_size):
+                buf[i] = data[i]
+            buf[payload_size] = 0
+            stream.codecpar.extradata = cython.cast(cython.p_uchar, buf)
+            stream.codecpar.extradata_size = payload_size
+
+        # Wrap as user-land stream.
+        meta_ptr = cython.address(stream.metadata)
+        err_check(lib.av_dict_set(meta_ptr, b"filename", name.encode(), 0))
+        mime_bytes = mimetype.encode()
+        err_check(lib.av_dict_set(meta_ptr, b"mimetype", mime_bytes, 0))
+
+        py_stream: Stream = wrap_stream(self, stream, None)
+        self.streams.add_stream(py_stream)
+        return py_stream
+
     def add_data_stream(self, codec_name=None, options: dict | None = None):
         """add_data_stream(codec_name=None)
 
@@ -270,21 +341,20 @@ def start_encoding(self):
         # Finalize and open all streams.
         for stream in self.streams:
             ctx = stream.codec_context
-            # Skip codec context handling for data streams without codecs
+            # Skip codec context handling for streams without codecs (e.g. data/attachments).
             if ctx is None:
-                if stream.type != "data":
+                if stream.type not in {"data", "attachment"}:
                     raise ValueError(f"Stream {stream.index} has no codec context")
-                continue
-
-            if not ctx.is_open:
-                for k, v in self.options.items():
-                    ctx.options.setdefault(k, v)
-                ctx.open()
-
-                # Track option consumption.
-                for k in self.options:
-                    if k not in ctx.options:
-                        used_options.add(k)
+            else:
+                if not ctx.is_open:
+                    for k, v in self.options.items():
+                        ctx.options.setdefault(k, v)
+                    ctx.open()
+
+                    # Track option consumption.
+                    for k in self.options:
+                        if k not in ctx.options:
+                            used_options.add(k)
 
             stream._finalize_for_output()
 

av/container/output.pyi

@@ -4,14 +4,14 @@ from typing import Sequence, TypeVar, Union, overload
 from av.audio import _AudioCodecName
 from av.audio.stream import AudioStream
 from av.packet import Packet
-from av.stream import DataStream
+from av.stream import AttachmentStream, DataStream, Stream
 from av.subtitles.stream import SubtitleStream
 from av.video import _VideoCodecName
 from av.video.stream import VideoStream
 
 from .core import Container
 
-_StreamT = TypeVar("_StreamT", bound=Union[VideoStream, AudioStream, SubtitleStream])
+_StreamT = TypeVar("_StreamT", bound=Stream)
 
 class OutputContainer(Container):
     def __enter__(self) -> OutputContainer: ...
@@ -42,6 +42,9 @@ class OutputContainer(Container):
     def add_stream_from_template(
         self, template: _StreamT, opaque: bool | None = None, **kwargs
     ) -> _StreamT: ...
+    def add_attachment(
+        self, name: str, mimetype: str, data: bytes
+    ) -> AttachmentStream: ...
     def add_data_stream(
         self, codec_name: str | None = None, options: dict[str, str] | None = None
     ) -> DataStream: ...

av/stream.py

@@ -148,6 +148,9 @@ def _finalize_for_output(self):
             errors=self.container.metadata_errors,
         )
 
+        if self.codec_context is None:
+            return
+
         if not self.ptr.time_base.num:
             self.ptr.time_base = self.codec_context.ptr.time_base
 
@@ -316,3 +319,17 @@ def mimetype(self):
         :rtype: str | None
         """
         return self.metadata.get("mimetype")
+
+    @property
+    def data(self):
+        """Return the raw attachment payload as bytes."""
+        extradata: cython.p_uchar = self.ptr.codecpar.extradata
+        size: cython.Py_ssize_t = self.ptr.codecpar.extradata_size
+        if extradata == cython.NULL or size <= 0:
+            return b""
+
+        payload = bytearray(size)
+        for i in range(size):
+            payload[i] = extradata[i]
+
+        return bytes(payload)

av/stream.pyi

@@ -59,3 +59,5 @@ class AttachmentStream(Stream):
     type: Literal["attachment"]
     @property
     def mimetype(self) -> str | None: ...
+    @property
+    def data(self) -> bytes: ...

tests/test_streams.py

@@ -1,10 +1,10 @@
 import os
 from fractions import Fraction
-from typing import Any, cast
 
 import pytest
 
 import av
+import av.datasets
 
 from .common import fate_suite
 
@@ -13,7 +13,14 @@ class TestStreams:
     @pytest.fixture(autouse=True)
     def cleanup(self):
         yield
-        for file in ("data.ts", "out.mkv"):
+        for file in (
+            "data.ts",
+            "data_source.ts",
+            "data_copy.ts",
+            "out.mkv",
+            "video_with_attachment.mkv",
+            "remuxed_attachment.mkv",
+        ):
             if os.path.exists(file):
                 os.remove(file)
 
@@ -149,21 +156,98 @@ def test_data_stream(self) -> None:
         container.close()
 
     def test_data_stream_from_template(self) -> None:
-        """Test that adding a data stream from a template raises ValueError."""
-
-        # Open an existing container with a data stream
-        input_container = av.open(fate_suite("mxf/track_01_v02.mxf"))
-        input_data_stream = input_container.streams.data[0]
-
-        # Create a new container and ensure using a data stream as a template raises ValueError
-        output_container = av.open("out.mkv", "w")
-        with pytest.raises(ValueError):
-            # input_data_stream is a DataStream at runtime; the test asserts that
-            # using it as a template raises ValueError. The static type stubs
-            # intentionally restrict which Stream subclasses are valid templates,
-            # so cast to Any here to keep the runtime check while satisfying
-            # the type checker.
-            output_container.add_stream_from_template(cast(Any, input_data_stream))
-
-        input_container.close()
-        output_container.close()
+        source_path = "data_source.ts"
+        payloads = [b"payload-a", b"payload-b", b"payload-c"]
+
+        with av.open(source_path, "w") as source:
+            source_stream = source.add_data_stream()
+            for i, payload in enumerate(payloads):
+                packet = av.Packet(payload)
+                packet.pts = i
+                packet.stream = source_stream
+                source.mux(packet)
+
+        copied_payloads: list[bytes] = []
+
+        with av.open(source_path) as input_container:
+            input_data_stream = input_container.streams.data[0]
+
+            with av.open("data_copy.ts", "w") as output_container:
+                output_data_stream = output_container.add_stream_from_template(
+                    input_data_stream
+                )
+
+                for packet in input_container.demux(input_data_stream):
+                    payload = bytes(packet)
+                    if not payload:
+                        continue
+                    copied_payloads.append(payload)
+                    clone = av.Packet(payload)
+                    clone.pts = packet.pts
+                    clone.dts = packet.dts
+                    clone.time_base = packet.time_base
+                    clone.stream = output_data_stream
+                    output_container.mux(clone)
+
+        with av.open("data_copy.ts") as remuxed:
+            output_stream = remuxed.streams.data[0]
+            assert output_stream.codec_context is None
+
+            remuxed_payloads: list[bytes] = []
+            for packet in remuxed.demux(output_stream):
+                payload = bytes(packet)
+                if payload:
+                    remuxed_payloads.append(payload)
+
+        assert remuxed_payloads == copied_payloads
+
+    def test_attachment_stream(self) -> None:
+        input_path = av.datasets.curated(
+            "pexels/time-lapse-video-of-night-sky-857195.mp4"
+        )
+        input_ = av.open(input_path)
+        out1_path = "video_with_attachment.mkv"
+
+        with av.open(out1_path, "w") as out1:
+            out1.add_attachment(
+                name="attachment.txt", mimetype="text/plain", data=b"hello\n"
+            )
+
+            in_v = input_.streams.video[0]
+            out_v = out1.add_stream_from_template(in_v)
+
+            for packet in input_.demux(in_v):
+                if packet.dts is None:
+                    continue
+                packet.stream = out_v
+                out1.mux(packet)
+
+        input_.close()
+
+        with av.open(out1_path) as c:
+            attachments = c.streams.attachments
+            assert len(attachments) == 1
+            att = attachments[0]
+            assert att.name == "attachment.txt"
+            assert att.mimetype == "text/plain"
+            assert att.data == b"hello\n"
+
+        out2_path = "remuxed_attachment.mkv"
+        with av.open(out1_path) as ic, av.open(out2_path, "w") as oc:
+            stream_map = {}
+            for s in ic.streams:
+                stream_map[s.index] = oc.add_stream_from_template(s)
+
+            for packet in ic.demux(ic.streams.video):
+                if packet.dts is None:
+                    continue
+                packet.stream = stream_map[packet.stream.index]
+                oc.mux(packet)
+
+        with av.open(out2_path) as c:
+            attachments = c.streams.attachments
+            assert len(attachments) == 1
+            att = attachments[0]
+            assert att.name == "attachment.txt"
+            assert att.mimetype == "text/plain"
+            assert att.data == b"hello\n"