diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a340bb088..3a69a86ab 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,13 @@ We are operating with `semantic versioning `_. Note that they these tags will not actually close the issue/PR until they are merged into the "default" branch. +v15.0.0 (Unreleased) +-------------------- + +Major: + +- Make ``SubtitleStream.decode()`` return the list of subtitles directly, without the intermediate ``SubtitleSet``. + v14.4.0 ------- diff --git a/av/subtitles/codeccontext.pxd b/av/subtitles/codeccontext.pxd index 42141aa4f..c94744e45 100644 --- a/av/subtitles/codeccontext.pxd +++ b/av/subtitles/codeccontext.pxd @@ -1,5 +1,6 @@ from av.codec.context cimport CodecContext +from av.packet cimport Packet cdef class SubtitleCodecContext(CodecContext): - pass + cpdef decode2(self, Packet packet) diff --git a/av/subtitles/codeccontext.py b/av/subtitles/codeccontext.py new file mode 100644 index 000000000..0b3dda063 --- /dev/null +++ b/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 diff --git a/av/subtitles/codeccontext.pyi b/av/subtitles/codeccontext.pyi index 0762c19f0..90c700935 100644 --- a/av/subtitles/codeccontext.pyi +++ b/av/subtitles/codeccontext.pyi @@ -1,6 +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: ... diff --git a/av/subtitles/codeccontext.pyx b/av/subtitles/codeccontext.pyx deleted file mode 100644 index c0712c92c..000000000 --- a/av/subtitles/codeccontext.pyx +++ /dev/null @@ -1,23 +0,0 @@ -cimport libav as lib - -from av.error cimport err_check -from av.packet cimport Packet -from av.subtitles.subtitle cimport SubtitleProxy, SubtitleSet - - -cdef class SubtitleCodecContext(CodecContext): - cdef _send_packet_and_recv(self, Packet packet): - if packet is None: - raise RuntimeError("packet cannot be None") - - cdef SubtitleProxy proxy = SubtitleProxy() - cdef int got_frame = 0 - - err_check( - lib.avcodec_decode_subtitle2(self.ptr, &proxy.struct, &got_frame, packet.ptr) - ) - - if got_frame: - return [SubtitleSet(proxy)] - else: - return [] diff --git a/av/subtitles/stream.py b/av/subtitles/stream.py new file mode 100644 index 000000000..525440e9d --- /dev/null +++ b/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) diff --git a/av/subtitles/stream.pyi b/av/subtitles/stream.pyi index cb1ac34a2..ac8083802 100644 --- a/av/subtitles/stream.pyi +++ b/av/subtitles/stream.pyi @@ -1,6 +1,9 @@ from av.packet import Packet from av.stream import Stream -from av.subtitles.subtitle import SubtitleSet +from av.subtitles.subtitle import AssSubtitle, BitmapSubtitle, SubtitleSet class SubtitleStream(Stream): - def decode(self, packet: Packet | None = None) -> list[SubtitleSet]: ... + def decode( + self, packet: Packet | None = None + ) -> list[AssSubtitle] | list[BitmapSubtitle]: ... + def decode2(self, packet: Packet) -> SubtitleSet | None: ... diff --git a/av/subtitles/stream.pyx b/av/subtitles/stream.pyx deleted file mode 100644 index 9f90b9871..000000000 --- a/av/subtitles/stream.pyx +++ /dev/null @@ -1,23 +0,0 @@ -from av.packet cimport Packet -from av.stream cimport Stream - - -cdef class SubtitleStream(Stream): - """ - A :class:`SubtitleStream` can contain many :class:`SubtitleSet` objects accessible via decoding. - """ - def __getattr__(self, name): - return getattr(self.codec_context, name) - - cpdef decode(self, Packet packet=None): - """ - Decode a :class:`.Packet` and return a list of :class:`.SubtitleSet`. - - :rtype: list[SubtitleSet] - - .. seealso:: This is a passthrough to :meth:`.CodecContext.decode`. - """ - if not packet: - packet = Packet() - - return self.codec_context.decode(packet) diff --git a/av/subtitles/subtitle.pyx b/av/subtitles/subtitle.py similarity index 51% rename from av/subtitles/subtitle.pyx rename to av/subtitles/subtitle.py index a713daa22..1acadf0b5 100644 --- a/av/subtitles/subtitle.pyx +++ b/av/subtitles/subtitle.py @@ -1,34 +1,49 @@ -from cpython cimport PyBuffer_FillInfo +import cython +from cython.cimports.cpython import PyBuffer_FillInfo, PyBytes_FromString +from cython.cimports.libc.stdint import uint64_t -cdef extern from "Python.h": - bytes PyBytes_FromString(char*) - - -cdef class SubtitleProxy: +@cython.cclass +class SubtitleProxy: def __dealloc__(self): - lib.avsubtitle_free(&self.struct) + lib.avsubtitle_free(cython.address(self.struct)) -cdef class SubtitleSet: +@cython.cclass +class SubtitleSet: """ A :class:`SubtitleSet` can contain many :class:`Subtitle` objects. + + Wraps :ffmpeg:`AVSubtitle`. """ - def __cinit__(self, SubtitleProxy proxy): + + def __cinit__(self, proxy: SubtitleProxy): self.proxy = proxy - self.rects = tuple(build_subtitle(self, i) for i in range(self.proxy.struct.num_rects)) + 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}>" + return ( + f"<{self.__class__.__module__}.{self.__class__.__name__} at 0x{id(self):x}>" + ) @property - def format(self): return self.proxy.struct.format + def format(self): + return self.proxy.struct.format + @property - def start_display_time(self): return self.proxy.struct.start_display_time + 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 + def end_display_time(self): + return self.proxy.struct.end_display_time + @property - def pts(self): return self.proxy.struct.pts + def pts(self): + """Same as packet pts, in av.time_base.""" + return self.proxy.struct.pts def __len__(self): return len(self.rects) @@ -40,33 +55,37 @@ def __getitem__(self, i): return self.rects[i] -cdef Subtitle build_subtitle(SubtitleSet subtitle, int index): +@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. - + The AVStream MUST be fully constructed and ready for use before this is called. """ - - if index < 0 or index >= subtitle.proxy.struct.num_rects: + if index < 0 or cython.cast(cython.uint, index) >= subtitle.proxy.struct.num_rects: raise ValueError("subtitle rect index out of range") - cdef lib.AVSubtitleRect *ptr = subtitle.proxy.struct.rects[index] + + ptr: cython.pointer[lib.AVSubtitleRect] = subtitle.proxy.struct.rects[index] if ptr.type == lib.SUBTITLE_BITMAP: return BitmapSubtitle(subtitle, index) - elif ptr.type == lib.SUBTITLE_ASS or ptr.type == lib.SUBTITLE_TEXT: + if ptr.type == lib.SUBTITLE_ASS or ptr.type == lib.SUBTITLE_TEXT: return AssSubtitle(subtitle, index) - else: - raise ValueError("unknown subtitle type %r" % ptr.type) + + raise ValueError("unknown subtitle type %r" % ptr.type) -cdef class Subtitle: +@cython.cclass +class Subtitle: """ An abstract base class for each concrete type of subtitle. Wraps :ffmpeg:`AVSubtitleRect` """ - def __cinit__(self, SubtitleSet subtitle, int index): - if index < 0 or index >= subtitle.proxy.struct.num_rects: + + 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] @@ -83,15 +102,14 @@ def __cinit__(self, SubtitleSet subtitle, int index): raise ValueError(f"unknown subtitle type {self.ptr.type!r}") def __repr__(self): - return f"<{self.__class__.__module__}.{self.__class__.__name__} at 0x{id(self):x}>" + return f"" -cdef class BitmapSubtitle(Subtitle): - def __cinit__(self, SubtitleSet subtitle, int index): +@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] + BitmapSubtitlePlane(self, i) for i in range(4) if self.ptr.linesize[i] ) def __repr__(self): @@ -101,15 +119,24 @@ def __repr__(self): ) @property - def x(self): return self.ptr.x + def x(self): + return self.ptr.x + @property - def y(self): return self.ptr.y + def y(self): + return self.ptr.y + @property - def width(self): return self.ptr.w + def width(self): + return self.ptr.w + @property - def height(self): return self.ptr.h + def height(self): + return self.ptr.h + @property - def nb_colors(self): return self.ptr.nb_colors + def nb_colors(self): + return self.ptr.nb_colors def __len__(self): return len(self.planes) @@ -121,8 +148,9 @@ def __getitem__(self, i): return self.planes[i] -cdef class BitmapSubtitlePlane: - def __cinit__(self, BitmapSubtitle subtitle, int index): +@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]: @@ -131,29 +159,28 @@ def __cinit__(self, BitmapSubtitle subtitle, int index): self.subtitle = subtitle self.index = index self.buffer_size = subtitle.ptr.w * subtitle.ptr.h - self._buffer = subtitle.ptr.data[index] + self._buffer = cython.cast(cython.p_void, subtitle.ptr.data[index]) # New-style buffer support. - def __getbuffer__(self, Py_buffer *view, int flags): + def __getbuffer__(self, view: cython.pointer[Py_buffer], flags: cython.int): PyBuffer_FillInfo(view, self, self._buffer, self.buffer_size, 0, flags) -cdef class AssSubtitle(Subtitle): +@cython.cclass +class AssSubtitle(Subtitle): """ Represents an ASS/Text subtitle format, as opposed to a bitmap Subtitle format. """ + def __repr__(self): - return ( - f"<{self.__class__.__module__}.{self.__class__.__name__} " - f"{self.text!r} at 0x{id(self):x}>" - ) + return f"" @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 NULL: + if self.ptr.ass is not cython.NULL: return PyBytes_FromString(self.ptr.ass) return b"" @@ -162,42 +189,44 @@ def dialogue(self): """ Extract the dialogue from the ass format. Strip comments. """ - comma_count = 0 - i = 0 - cdef bytes ass_text = self.ass - cdef bytes result = b"" - - while comma_count < 8 and i < len(ass_text): - if bytes([ass_text[i]]) == b",": + 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 - state = False - while i < len(ass_text): - char = bytes([ass_text[i]]) - next_char = b"" if i + 1 >= len(ass_text) else bytes([ass_text[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"\\" and next_char == b"N": - result += b"\n" + if char == b"\\"[0] and next_char == b"N"[0]: + result.append(b"\n"[0]) i += 2 continue if not state: - if char == b"{" and next_char != b"\\": + if char == b"{"[0] and next_char != b"\\"[0]: state = True else: - result += char - elif char == b"}": + result.append(char) + elif char == b"}"[0]: state = False i += 1 - return result + return bytes(result) @property def text(self): """ Rarely used attribute. You're probably looking for dialogue. """ - if self.ptr.text is not NULL: + if self.ptr.text is not cython.NULL: return PyBytes_FromString(self.ptr.text) return b"" diff --git a/tests/test_subtitles.py b/tests/test_subtitles.py index e2b5ab512..66e6f9cfa 100644 --- a/tests/test_subtitles.py +++ b/tests/test_subtitles.py @@ -1,3 +1,5 @@ +from typing import cast + import av from av.subtitles.subtitle import AssSubtitle, BitmapSubtitle @@ -8,43 +10,46 @@ class TestSubtitle: def test_movtext(self) -> None: path = fate_suite("sub/MovText_capability_tester.mp4") - subs = [] + subs: list[AssSubtitle] = [] with av.open(path) as container: for packet in container.demux(): - subs.extend(packet.decode()) + subs.extend(cast(list[AssSubtitle], packet.decode())) assert len(subs) == 3 - subset = subs[0] - assert subset.format == 1 - assert subset.pts == 970000 - assert subset.start_display_time == 0 - assert subset.end_display_time == 1570 - - sub = subset[0] + sub = subs[0] assert isinstance(sub, AssSubtitle) assert sub.type == b"ass" assert sub.text == b"" assert sub.ass == b"0,0,Default,,0,0,0,,- Test 1.\\N- Test 2." assert sub.dialogue == b"- Test 1.\n- Test 2." + def test_subset(self) -> None: + path = fate_suite("sub/MovText_capability_tester.mp4") + + with av.open(path) as container: + subs = container.streams.subtitles[0] + for packet in container.demux(subs): + subset = subs.decode2(packet) + if subset is not None: + assert not isinstance(subset, av.subtitles.subtitle.Subtitle) + assert isinstance(subset, av.subtitles.subtitle.SubtitleSet) + assert subset.format == 1 + assert hasattr(subset, "pts") + assert subset.start_display_time == 0 + assert hasattr(subset, "end_display_time") + def test_vobsub(self) -> None: path = fate_suite("sub/vobsub.sub") - subs = [] + subs: list[BitmapSubtitle] = [] with av.open(path) as container: for packet in container.demux(): - subs.extend(packet.decode()) + subs.extend(cast(list[BitmapSubtitle], packet.decode())) assert len(subs) == 43 - subset = subs[0] - assert subset.format == 0 - assert subset.pts == 132499044 - assert subset.start_display_time == 0 - assert subset.end_display_time == 4960 - - sub = subset[0] + sub = subs[0] assert isinstance(sub, BitmapSubtitle) assert sub.type == b"bitmap" assert sub.x == 259 @@ -60,7 +65,7 @@ def test_vobsub(self) -> None: def test_subtitle_flush(self) -> None: path = fate_suite("sub/MovText_capability_tester.mp4") - subs = [] + subs: list[object] = [] with av.open(path) as container: stream = container.streams.subtitles[0] for packet in container.demux(stream):