cpu-diffusers.upgrade-tmp/venv/lib/python3.10/site-packages/av/subtitles/codeccontext.py

import cython
from cython.cimports import libav as lib
from cython.cimports.av.buffer import ByteSource, bytesource
from cython.cimports.av.codec.context import CodecContext
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.ccall
    def decode(self, packet: Packet | None = None):
        """Decode a subtitle packet, returning a list of :class:`.Subtitle` objects
        if a subtitle was decoded, or an empty list otherwise."""
        if not self.codec.ptr:
            raise ValueError("cannot decode unknown codec")

        if packet is None:
            raise RuntimeError("packet cannot be None")

        self.open(strict=False)
        proxy: SubtitleProxy = SubtitleProxy()
        got_frame: cython.int = 0

        err_check(
            lib.avcodec_decode_subtitle2(
                self.ptr,
                cython.address(proxy.struct),
                cython.address(got_frame),
                packet.ptr,
            )
        )

        if got_frame:
            return list(SubtitleSet(proxy))
        return []

    @cython.ccall
    def decode2(self, packet: Packet):
        """
        Returns SubtitleSet if you really need it.
        """
        if not self.codec.ptr:
            raise ValueError("cannot decode unknown codec")

        self.open(strict=False)

        proxy: SubtitleProxy = SubtitleProxy()
        got_frame: cython.int = 0

        err_check(
            lib.avcodec_decode_subtitle2(
                self.ptr,
                cython.address(proxy.struct),
                cython.address(got_frame),
                packet.ptr,
            )
        )

        if got_frame:
            return SubtitleSet(proxy)
        return None