add read me

venv/lib/python3.12/site-packages/av/subtitles/codeccontext.pxd

@@ -0,0 +1,6 @@
+from av.codec.context cimport CodecContext
+from av.packet cimport Packet
+
+
+cdef class SubtitleCodecContext(CodecContext):
+    cpdef decode2(self, Packet packet)

venv/lib/python3.12/site-packages/av/subtitles/codeccontext.py

@@ -0,0 +1,55 @@
+import cython
+from cython.cimports import libav as lib
+from cython.cimports.av.error import err_check
+from cython.cimports.av.packet import Packet
+from cython.cimports.av.subtitles.subtitle import SubtitleProxy, SubtitleSet
+
+
+@cython.cclass
+class SubtitleCodecContext(CodecContext):
+    @cython.cfunc
+    def _send_packet_and_recv(self, packet: Packet | None):
+        if packet is None:
+            raise RuntimeError("packet cannot be None")
+
+        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 []
+
+    @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

venv/lib/python3.12/site-packages/av/subtitles/codeccontext.pyi

@@ -0,0 +1,9 @@
+from typing import Literal
+
+from av.codec.context import CodecContext
+from av.packet import Packet
+from av.subtitles.subtitle import SubtitleSet
+
+class SubtitleCodecContext(CodecContext):
+    type: Literal["subtitle"]
+    def decode2(self, packet: Packet) -> SubtitleSet | None: ...

venv/lib/python3.12/site-packages/av/subtitles/stream.pxd

@@ -0,0 +1,6 @@
+from av.packet cimport Packet
+from av.stream cimport Stream
+
+
+cdef class SubtitleStream(Stream):
+    cpdef decode(self, Packet packet=?)

venv/lib/python3.12/site-packages/av/subtitles/stream.py

@@ -0,0 +1,23 @@
+import cython
+from cython.cimports.av.packet import Packet
+from cython.cimports.av.stream import Stream
+
+
+@cython.cclass
+class SubtitleStream(Stream):
+    def __getattr__(self, name):
+        return getattr(self.codec_context, name)
+
+    @cython.ccall
+    def decode(self, packet: Packet | None = None):
+        """
+        Decode a :class:`.Packet` and returns a subtitle object.
+
+        :rtype: list[AssSubtitle] | list[BitmapSubtitle]
+
+        .. seealso:: This is a passthrough to :meth:`.CodecContext.decode`.
+        """
+        if not packet:
+            packet = Packet()
+
+        return self.codec_context.decode(packet)

venv/lib/python3.12/site-packages/av/subtitles/stream.pyi

@@ -0,0 +1,9 @@
+from av.packet import Packet
+from av.stream import Stream
+from av.subtitles.subtitle import AssSubtitle, BitmapSubtitle, SubtitleSet
+
+class SubtitleStream(Stream):
+    def decode(
+        self, packet: Packet | None = None
+    ) -> list[AssSubtitle] | list[BitmapSubtitle]: ...
+    def decode2(self, packet: Packet) -> SubtitleSet | None: ...

venv/lib/python3.12/site-packages/av/subtitles/subtitle.pxd

@@ -0,0 +1,31 @@
+cimport libav as lib
+
+
+cdef class SubtitleProxy:
+    cdef lib.AVSubtitle struct
+
+
+cdef class SubtitleSet:
+    cdef SubtitleProxy proxy
+    cdef readonly tuple rects
+
+
+cdef class Subtitle:
+    cdef SubtitleProxy proxy
+    cdef lib.AVSubtitleRect *ptr
+    cdef readonly bytes type
+
+cdef class TextSubtitle(Subtitle):
+    pass
+
+cdef class ASSSubtitle(Subtitle):
+    pass
+
+cdef class BitmapSubtitle(Subtitle):
+    cdef readonly planes
+
+cdef class BitmapSubtitlePlane:
+    cdef readonly BitmapSubtitle subtitle
+    cdef readonly int index
+    cdef readonly long buffer_size
+    cdef void *_buffer

venv/lib/python3.12/site-packages/av/subtitles/subtitle.py

@@ -0,0 +1,232 @@
+import cython
+from cython.cimports.cpython import PyBuffer_FillInfo, PyBytes_FromString
+from cython.cimports.libc.stdint import uint64_t
+
+
+@cython.cclass
+class SubtitleProxy:
+    def __dealloc__(self):
+        lib.avsubtitle_free(cython.address(self.struct))
+
+
+@cython.cclass
+class SubtitleSet:
+    """
+    A :class:`SubtitleSet` can contain many :class:`Subtitle` objects.
+
+    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 __repr__(self):
+        return (
+            f"<{self.__class__.__module__}.{self.__class__.__name__} at 0x{id(self):x}>"
+        )
+
+    @property
+    def format(self):
+        return self.proxy.struct.format
+
+    @property
+    def start_display_time(self):
+        return self.proxy.struct.start_display_time
+
+    @property
+    def end_display_time(self):
+        return self.proxy.struct.end_display_time
+
+    @property
+    def pts(self):
+        """Same as packet pts, in av.time_base."""
+        return self.proxy.struct.pts
+
+    def __len__(self):
+        return len(self.rects)
+
+    def __iter__(self):
+        return iter(self.rects)
+
+    def __getitem__(self, i):
+        return self.rects[i]
+
+
+@cython.cfunc
+def build_subtitle(subtitle: SubtitleSet, index: cython.int) -> Subtitle:
+    """Build an av.Stream for an existing AVStream.
+
+    The AVStream MUST be fully constructed and ready for use before this is called.
+    """
+    if index < 0 or cython.cast(cython.uint, index) >= subtitle.proxy.struct.num_rects:
+        raise ValueError("subtitle rect index out of range")
+
+    ptr: cython.pointer[lib.AVSubtitleRect] = subtitle.proxy.struct.rects[index]
+
+    if ptr.type == lib.SUBTITLE_BITMAP:
+        return BitmapSubtitle(subtitle, index)
+    if ptr.type == lib.SUBTITLE_ASS or ptr.type == lib.SUBTITLE_TEXT:
+        return AssSubtitle(subtitle, index)
+
+    raise ValueError("unknown subtitle type %r" % ptr.type)
+
+
+@cython.cclass
+class Subtitle:
+    """
+    An abstract base class for each concrete type of subtitle.
+    Wraps :ffmpeg:`AVSubtitleRect`
+    """
+
+    def __cinit__(self, subtitle: SubtitleSet, index: cython.int):
+        if (
+            index < 0
+            or cython.cast(cython.uint, index) >= subtitle.proxy.struct.num_rects
+        ):
+            raise ValueError("subtitle rect index out of range")
+        self.proxy = subtitle.proxy
+        self.ptr = self.proxy.struct.rects[index]
+
+        if self.ptr.type == lib.SUBTITLE_NONE:
+            self.type = b"none"
+        elif self.ptr.type == lib.SUBTITLE_BITMAP:
+            self.type = b"bitmap"
+        elif self.ptr.type == lib.SUBTITLE_TEXT:
+            self.type = b"text"
+        elif self.ptr.type == lib.SUBTITLE_ASS:
+            self.type = b"ass"
+        else:
+            raise ValueError(f"unknown subtitle type {self.ptr.type!r}")
+
+    def __repr__(self):
+        return f"<av.{self.__class__.__name__} at 0x{id(self):x}>"
+
+
+@cython.cclass
+class BitmapSubtitle(Subtitle):
+    def __cinit__(self, subtitle: SubtitleSet, index: cython.int):
+        self.planes = tuple(
+            BitmapSubtitlePlane(self, i) for i in range(4) if self.ptr.linesize[i]
+        )
+
+    def __repr__(self):
+        return (
+            f"<{self.__class__.__module__}.{self.__class__.__name__} "
+            f"{self.width}x{self.height} at {self.x},{self.y}; at 0x{id(self):x}>"
+        )
+
+    @property
+    def x(self):
+        return self.ptr.x
+
+    @property
+    def y(self):
+        return self.ptr.y
+
+    @property
+    def width(self):
+        return self.ptr.w
+
+    @property
+    def height(self):
+        return self.ptr.h
+
+    @property
+    def nb_colors(self):
+        return self.ptr.nb_colors
+
+    def __len__(self):
+        return len(self.planes)
+
+    def __iter__(self):
+        return iter(self.planes)
+
+    def __getitem__(self, i):
+        return self.planes[i]
+
+
+@cython.cclass
+class BitmapSubtitlePlane:
+    def __cinit__(self, subtitle: BitmapSubtitle, index: cython.int):
+        if index >= 4:
+            raise ValueError("BitmapSubtitles have only 4 planes")
+        if not subtitle.ptr.linesize[index]:
+            raise ValueError("plane does not exist")
+
+        self.subtitle = subtitle
+        self.index = index
+        self.buffer_size = subtitle.ptr.w * subtitle.ptr.h
+        self._buffer = cython.cast(cython.p_void, subtitle.ptr.data[index])
+
+    # New-style buffer support.
+    def __getbuffer__(self, view: cython.pointer[Py_buffer], flags: cython.int):
+        PyBuffer_FillInfo(view, self, self._buffer, self.buffer_size, 0, flags)
+
+
+@cython.cclass
+class AssSubtitle(Subtitle):
+    """
+    Represents an ASS/Text subtitle format, as opposed to a bitmap Subtitle format.
+    """
+
+    def __repr__(self):
+        return f"<av.AssSubtitle {self.dialogue!r} at 0x{id(self):x}>"
+
+    @property
+    def ass(self):
+        """
+        Returns the subtitle in the ASS/SSA format. Used by the vast majority of subtitle formats.
+        """
+        if self.ptr.ass is not cython.NULL:
+            return PyBytes_FromString(self.ptr.ass)
+        return b""
+
+    @property
+    def dialogue(self):
+        """
+        Extract the dialogue from the ass format. Strip comments.
+        """
+        comma_count: cython.short = 0
+        i: uint64_t = 0
+        state: cython.bint = False
+        ass_text: bytes = self.ass
+        char, next_char = cython.declare(cython.char)
+        result: bytearray = bytearray()
+        text_len: cython.Py_ssize_t = len(ass_text)
+
+        while comma_count < 8 and i < text_len:
+            if ass_text[i] == b","[0]:
+                comma_count += 1
+            i += 1
+
+        while i < text_len:
+            char = ass_text[i]
+            next_char = 0 if i + 1 >= text_len else ass_text[i + 1]
+
+            if char == b"\\"[0] and next_char == b"N"[0]:
+                result.append(b"\n"[0])
+                i += 2
+                continue
+
+            if not state:
+                if char == b"{"[0] and next_char != b"\\"[0]:
+                    state = True
+                else:
+                    result.append(char)
+            elif char == b"}"[0]:
+                state = False
+            i += 1
+
+        return bytes(result)
+
+    @property
+    def text(self):
+        """
+        Rarely used attribute. You're probably looking for dialogue.
+        """
+        if self.ptr.text is not cython.NULL:
+            return PyBytes_FromString(self.ptr.text)
+        return b""

venv/lib/python3.12/site-packages/av/subtitles/subtitle.pyi

@@ -0,0 +1,37 @@
+from typing import Iterator, Literal
+
+class SubtitleSet:
+    format: int
+    start_display_time: int
+    end_display_time: int
+    pts: int
+    rects: tuple[Subtitle]
+
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[Subtitle]: ...
+    def __getitem__(self, i: int) -> Subtitle: ...
+
+class Subtitle: ...
+
+class BitmapSubtitle(Subtitle):
+    type: Literal[b"bitmap"]
+    x: int
+    y: int
+    width: int
+    height: int
+    nb_colors: int
+    planes: tuple[BitmapSubtitlePlane, ...]
+
+class BitmapSubtitlePlane:
+    subtitle: BitmapSubtitle
+    index: int
+    buffer_size: int
+
+class AssSubtitle(Subtitle):
+    type: Literal[b"ass", b"text"]
+    @property
+    def ass(self) -> bytes: ...
+    @property
+    def dialogue(self) -> bytes: ...
+    @property
+    def text(self) -> bytes: ...