from cpython cimport PyBuffer_FillInfo
cdef extern from "Python.h":
bytes PyBytes_FromString(char*)
cdef class SubtitleProxy:
def __dealloc__(self):
lib.avsubtitle_free(&self.struct)
cdef class SubtitleSet:
"""
A :class:`SubtitleSet` can contain many :class:`Subtitle` objects.
"""
def __cinit__(self, SubtitleProxy proxy):
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): 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]
cdef Subtitle build_subtitle(SubtitleSet subtitle, int index):
"""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 <unsigned int>index >= subtitle.proxy.struct.num_rects:
raise ValueError("subtitle rect index out of range")
cdef lib.AVSubtitleRect *ptr = 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:
return AssSubtitle(subtitle, index)
else:
raise ValueError("unknown subtitle type %r" % ptr.type)
cdef 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 <unsigned int>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"<{self.__class__.__module__}.{self.__class__.__name__} at 0x{id(self):x}>"
cdef class BitmapSubtitle(Subtitle):
def __cinit__(self, SubtitleSet subtitle, int index):
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]
cdef class BitmapSubtitlePlane:
def __cinit__(self, BitmapSubtitle subtitle, int index):
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 = <void*>subtitle.ptr.data[index]
# New-style buffer support.
def __getbuffer__(self, Py_buffer *view, int flags):
PyBuffer_FillInfo(view, self, self._buffer, self.buffer_size, 0, flags)
cdef 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}>"
)
@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:
return PyBytes_FromString(self.ptr.ass)
return b""
@property
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 += 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]])
if char == b"\\" and next_char == b"N":
result += b"\n"
i += 2
continue
if not state:
if char == b"{" and next_char != b"\\":
state = True
else:
result += char
elif char == b"}":
state = False
i += 1
return result
@property
def text(self):
"""
Rarely used attribute. You're probably looking for dialogue.
"""
if self.ptr.text is not NULL:
return PyBytes_FromString(self.ptr.text)
return b""