diff --git a/av/container/core.pyi b/av/container/core.pyi index 8cd2a9dc5..f923ba01e 100644 --- a/av/container/core.pyi +++ b/av/container/core.pyi @@ -67,7 +67,7 @@ class AudioCodec(IntEnum): pcm_u8 = cast(int, ...) pcm_vidc = cast(int, ...) -class _Chapter(TypedDict): +class Chapter(TypedDict): id: int start: int end: int @@ -102,7 +102,8 @@ class Container: ) -> bool: ... def set_timeout(self, timeout: Real | None) -> None: ... def start_timeout(self) -> None: ... - def chapters(self) -> list[_Chapter]: ... + def chapters(self) -> list[Chapter]: ... + def set_chapters(self, chapters: list[Chapter]) -> None: ... @overload def open( diff --git a/av/container/core.pyx b/av/container/core.pyx index 076faea15..299ac6ed3 100755 --- a/av/container/core.pyx +++ b/av/container/core.pyx @@ -15,7 +15,12 @@ from av.container.output cimport OutputContainer from av.container.pyio cimport pyio_close_custom_gil, pyio_close_gil from av.error cimport err_check, stash_exception from av.format cimport build_container_format -from av.utils cimport avdict_to_dict, avrational_to_fraction +from av.utils cimport ( + avdict_to_dict, + avrational_to_fraction, + dict_to_avdict, + to_avrational, +) from av.dictionary import Dictionary from av.logging import Capture as LogCapture @@ -123,6 +128,17 @@ cdef int pyav_io_close_gil(lib.AVFormatContext *s, lib.AVIOContext *pb) noexcept return result +cdef void _free_chapters(lib.AVFormatContext *ctx) noexcept nogil: + cdef int i + if ctx.chapters != NULL: + for i in range(ctx.nb_chapters): + if ctx.chapters[i] != NULL: + if ctx.chapters[i].metadata != NULL: + lib.av_dict_free(&ctx.chapters[i].metadata) + lib.av_freep(&ctx.chapters[i]) + lib.av_freep(&ctx.chapters) + ctx.nb_chapters = 0 + class Flags(Flag): gen_pts: "Generate missing pts even if it requires parsing future frames." = lib.AVFMT_FLAG_GENPTS @@ -346,6 +362,39 @@ cdef class Container: }) return result + def set_chapters(self, chapters): + self._assert_open() + + cdef int count = len(chapters) + cdef int i + cdef lib.AVChapter **ch_array + cdef lib.AVChapter *ch + cdef dict entry + + with nogil: + _free_chapters(self.ptr) + + ch_array = lib.av_malloc(count * sizeof(lib.AVChapter *)) + if ch_array == NULL: + raise MemoryError("av_malloc failed for chapters") + + for i in range(count): + entry = chapters[i] + ch = lib.av_malloc(sizeof(lib.AVChapter)) + if ch == NULL: + raise MemoryError("av_malloc failed for chapter") + ch.id = entry["id"] + ch.start = entry["start"] + ch.end = entry["end"] + to_avrational(entry["time_base"], &ch.time_base) + ch.metadata = NULL + if "metadata" in entry: + dict_to_avdict(&ch.metadata, entry["metadata"], self.metadata_encoding, self.metadata_errors) + ch_array[i] = ch + + self.ptr.nb_chapters = count + self.ptr.chapters = ch_array + def open( file, mode=None, diff --git a/tests/test_chapters.py b/tests/test_chapters.py index 6a3d371b2..c8a3626ce 100644 --- a/tests/test_chapters.py +++ b/tests/test_chapters.py @@ -39,3 +39,20 @@ def test_chapters() -> None: path = fate_suite("vorbis/vorbis_chapter_extension_demo.ogg") with av.open(path) as container: assert container.chapters() == expected + + +def test_set_chapters() -> None: + chapters: list[av.container.Chapter] = [ + { + "id": 1, + "start": 0, + "end": 5000, + "time_base": Fraction(1, 1000), + "metadata": {"title": "start"}, + } + ] + + path = fate_suite("h264/interlaced_crop.mp4") + with av.open(path) as container: + container.set_chapters(chapters) + assert container.chapters() == chapters