Add subtitle encoding support

Author: skeskinen

av/subtitles/codeccontext.pxd

@@ -3,4 +3,5 @@ from av.packet cimport Packet
 
 
 cdef class SubtitleCodecContext(CodecContext):
+    cdef bint subtitle_header_set
     cpdef decode2(self, Packet packet)

av/subtitles/codeccontext.py

@@ -1,12 +1,108 @@
 import cython
 from cython.cimports import libav as lib
+from cython.cimports.av.bytesource import ByteSource, bytesource
 from cython.cimports.av.error import err_check
 from cython.cimports.av.packet import Packet
 from cython.cimports.av.subtitles.subtitle import SubtitleProxy, SubtitleSet
+from cython.cimports.cpython.bytes import PyBytes_FromStringAndSize
+from cython.cimports.libc.string import memcpy, strlen
 
 
 @cython.cclass
 class SubtitleCodecContext(CodecContext):
+    @property
+    def subtitle_header(self) -> bytes | None:
+        """Get the subtitle header data (ASS/SSA format for text subtitles)."""
+        if (
+            self.ptr.subtitle_header == cython.NULL
+            or self.ptr.subtitle_header_size <= 0
+        ):
+            return None
+        return PyBytes_FromStringAndSize(
+            cython.cast(cython.p_char, self.ptr.subtitle_header),
+            self.ptr.subtitle_header_size,
+        )
+
+    @subtitle_header.setter
+    def subtitle_header(self, data: bytes | None) -> None:
+        """Set the subtitle header data."""
+        source: ByteSource
+        if data is None:
+            lib.av_freep(cython.address(self.ptr.subtitle_header))
+            self.ptr.subtitle_header_size = 0
+        else:
+            source = bytesource(data)
+            self.ptr.subtitle_header = cython.cast(
+                cython.p_uchar,
+                lib.av_realloc(
+                    self.ptr.subtitle_header,
+                    source.length + lib.AV_INPUT_BUFFER_PADDING_SIZE,
+                ),
+            )
+            if not self.ptr.subtitle_header:
+                raise MemoryError("Cannot allocate subtitle_header")
+            memcpy(self.ptr.subtitle_header, source.ptr, source.length)
+            self.ptr.subtitle_header_size = source.length
+        self.subtitle_header_set = True
+
+    def __dealloc__(self) -> None:
+        if self.ptr and self.subtitle_header_set:
+            lib.av_freep(cython.address(self.ptr.subtitle_header))
+
+    def encode_subtitle(self, subtitle: SubtitleSet) -> Packet:
+        """
+        Encode a SubtitleSet into a Packet.
+
+        Args:
+            subtitle: The SubtitleSet to encode
+
+        Returns:
+            A Packet containing the encoded subtitle data
+        """
+        if not self.codec.ptr:
+            raise ValueError("Cannot encode with unknown codec")
+
+        self.open(strict=False)
+
+        # Calculate buffer size from subtitle text length
+        buf_size: cython.size_t = 0
+        i: cython.uint
+        for i in range(subtitle.proxy.struct.num_rects):
+            rect = subtitle.proxy.struct.rects[i]
+            if rect.ass != cython.NULL:
+                buf_size += strlen(rect.ass)
+            if rect.text != cython.NULL:
+                buf_size += strlen(rect.text)
+        buf_size += 1024  # padding for format overhead
+
+        buf: cython.p_uchar = cython.cast(cython.p_uchar, lib.av_malloc(buf_size))
+        if buf == cython.NULL:
+            raise MemoryError("Failed to allocate subtitle encode buffer")
+
+        ret: cython.int = lib.avcodec_encode_subtitle(
+            self.ptr,
+            buf,
+            buf_size,
+            cython.address(subtitle.proxy.struct),
+        )
+
+        if ret < 0:
+            lib.av_free(buf)
+            err_check(ret, "avcodec_encode_subtitle()")
+
+        packet: Packet = Packet(ret)
+        memcpy(packet.ptr.data, buf, ret)
+        lib.av_free(buf)
+
+        packet.ptr.pts = subtitle.proxy.struct.pts
+        packet.ptr.dts = subtitle.proxy.struct.pts
+        packet.ptr.duration = (
+            subtitle.proxy.struct.end_display_time
+            - subtitle.proxy.struct.start_display_time
+        )
+
+        return packet
+
     @cython.cfunc
     def _send_packet_and_recv(self, packet: Packet | None):
         if packet is None:

av/subtitles/codeccontext.pyi

@@ -6,4 +6,6 @@ from av.subtitles.subtitle import SubtitleSet
 
 class SubtitleCodecContext(CodecContext):
     type: Literal["subtitle"]
+    subtitle_header: bytes | None
     def decode2(self, packet: Packet) -> SubtitleSet | None: ...
+    def encode_subtitle(self, subtitle: SubtitleSet) -> Packet: ...

av/subtitles/subtitle.py

@@ -1,6 +1,7 @@
 import cython
 from cython.cimports.cpython import PyBuffer_FillInfo, PyBytes_FromString
-from cython.cimports.libc.stdint import uint64_t
+from cython.cimports.libc.stdint import int64_t, uint64_t
+from cython.cimports.libc.string import memcpy, strlen
 
 
 @cython.cclass
@@ -9,6 +10,9 @@ def __dealloc__(self):
         lib.avsubtitle_free(cython.address(self.struct))
 
 
+_cinit_bypass_sentinel = object()
+
+
 @cython.cclass
 class SubtitleSet:
     """
@@ -17,11 +21,94 @@ class SubtitleSet:
     Wraps :ffmpeg:`AVSubtitle`.
     """
 
-    def __cinit__(self, proxy: SubtitleProxy):
-        self.proxy = proxy
-        self.rects = tuple(
-            build_subtitle(self, i) for i in range(self.proxy.struct.num_rects)
+    def __cinit__(self, proxy_or_sentinel=None):
+        if proxy_or_sentinel is _cinit_bypass_sentinel:
+            # Creating empty SubtitleSet for encoding
+            self.proxy = SubtitleProxy()
+            self.rects = ()
+        elif isinstance(proxy_or_sentinel, SubtitleProxy):
+            # Creating from decoded subtitle
+            self.proxy = proxy_or_sentinel
+            self.rects = tuple(
+                build_subtitle(self, i) for i in range(self.proxy.struct.num_rects)
+            )
+        else:
+            raise TypeError(
+                "SubtitleSet requires a SubtitleProxy or use SubtitleSet.create()"
+            )
+
+    @staticmethod
+    def create(
+        text: bytes,
+        start: int,
+        end: int,
+        pts: int = 0,
+        subtitle_format: int = 1,
+    ) -> "SubtitleSet":
+        """
+        Create a SubtitleSet for encoding.
+
+        Args:
+            text: The subtitle text in ASS dialogue format
+                  (e.g. b"0,0,Default,,0,0,0,,Hello World")
+            start: Start display time as offset from pts (typically 0)
+            end: End display time as offset from pts (i.e., duration)
+            pts: Presentation timestamp in stream time_base units
+            subtitle_format: Subtitle format (default 1 for text)
+
+        Note:
+            All timing values should be in stream time_base units.
+            For MKV (time_base=1/1000), units are milliseconds.
+            For MP4 (time_base=1/1000000), units are microseconds.
+
+        Returns:
+            A SubtitleSet ready for encoding
+        """
+        subset: SubtitleSet = SubtitleSet(_cinit_bypass_sentinel)
+
+        subset.proxy.struct.format = subtitle_format
+        subset.proxy.struct.start_display_time = start
+        subset.proxy.struct.end_display_time = end
+        subset.proxy.struct.pts = pts
+
+        subset.proxy.struct.num_rects = 1
+        subset.proxy.struct.rects = cython.cast(
+            cython.pointer[cython.pointer[lib.AVSubtitleRect]],
+            lib.av_mallocz(cython.sizeof(cython.pointer[lib.AVSubtitleRect])),
+        )
+        if subset.proxy.struct.rects == cython.NULL:
+            raise MemoryError("Failed to allocate subtitle rects array")
+
+        rect: cython.pointer[lib.AVSubtitleRect] = cython.cast(
+            cython.pointer[lib.AVSubtitleRect],
+            lib.av_mallocz(cython.sizeof(lib.AVSubtitleRect)),
         )
+        if rect == cython.NULL:
+            lib.av_free(subset.proxy.struct.rects)
+            subset.proxy.struct.rects = cython.NULL
+            raise MemoryError("Failed to allocate subtitle rect")
+
+        subset.proxy.struct.rects[0] = rect
+
+        rect.x = 0
+        rect.y = 0
+        rect.w = 0
+        rect.h = 0
+        rect.nb_colors = 0
+        rect.type = lib.SUBTITLE_ASS
+        rect.text = cython.NULL
+        rect.flags = 0
+
+        text_len: cython.Py_ssize_t = len(text)
+        rect.ass = cython.cast(cython.p_char, lib.av_malloc(text_len + 1))
+        if rect.ass == cython.NULL:
+            raise MemoryError("Failed to allocate subtitle text")
+        memcpy(rect.ass, cython.cast(cython.p_char, text), text_len)
+        rect.ass[text_len] = 0
+
+        subset.rects = (AssSubtitle(subset, 0),)
+
+        return subset
 
     def __repr__(self):
         return (
@@ -32,19 +119,35 @@ def __repr__(self):
     def format(self):
         return self.proxy.struct.format
 
+    @format.setter
+    def format(self, value: int):
+        self.proxy.struct.format = value
+
     @property
     def start_display_time(self):
         return self.proxy.struct.start_display_time
 
+    @start_display_time.setter
+    def start_display_time(self, value: int):
+        self.proxy.struct.start_display_time = value
+
     @property
     def end_display_time(self):
         return self.proxy.struct.end_display_time
 
+    @end_display_time.setter
+    def end_display_time(self, value: int):
+        self.proxy.struct.end_display_time = value
+
     @property
     def pts(self):
         """Same as packet pts, in av.time_base."""
         return self.proxy.struct.pts
 
+    @pts.setter
+    def pts(self, value: int):
+        self.proxy.struct.pts = value
+
     def __len__(self):
         return len(self.rects)
 

av/subtitles/subtitle.pyi

@@ -5,8 +5,16 @@ class SubtitleSet:
     start_display_time: int
     end_display_time: int
     pts: int
-    rects: tuple[Subtitle]
+    rects: tuple[Subtitle, ...]
 
+    @staticmethod
+    def create(
+        text: bytes,
+        start: int,
+        end: int,
+        pts: int = 0,
+        subtitle_format: int = 1,
+    ) -> SubtitleSet: ...
     def __len__(self) -> int: ...
     def __iter__(self) -> Iterator[Subtitle]: ...
     def __getitem__(self, i: int) -> Subtitle: ...

include/libavcodec/avcodec.pxd

@@ -289,6 +289,10 @@ cdef extern from "libavcodec/avcodec.h" nogil:
         int extradata_size
         uint8_t *extradata
 
+        # Subtitle header (ASS/SSA format for text subtitles)
+        uint8_t *subtitle_header
+        int subtitle_header_size
+
         int delay
 
         AVCodec *codec

include/libavutil/avutil.pxd

@@ -110,9 +110,11 @@ cdef extern from "libavutil/avutil.h" nogil:
     cdef double M_PI
 
     cdef void* av_malloc(size_t size)
+    cdef void* av_mallocz(size_t size)
     cdef void *av_calloc(size_t nmemb, size_t size)
     cdef void *av_realloc(void *ptr, size_t size)
 
+    cdef void av_free(void *ptr)
     cdef void av_freep(void *ptr)
 
     cdef int av_get_bytes_per_sample(AVSampleFormat sample_fmt)