Implement set_chapters method

Author: DE-AI

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(

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(<void **>&ctx.chapters[i])
+            lib.av_freep(<void **>&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.AVChapter **>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.AVChapter *>lib.av_malloc(sizeof(lib.AVChapter))
+            if ch == NULL:
+                raise MemoryError("av_malloc failed for chapter")
+            ch.id = entry["id"]
+            ch.start = <int64_t>entry["start"]
+            ch.end = <int64_t>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,

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