Expose AVIndexEntry (#2136)

Queuecumber · web-flow · commit f462b75df519 · 2026-01-27T18:32:59.000-05:00
diff --git a/av/indexentries.pxd b/av/indexentries.pxd
@@ -0,0 +1,9 @@
+cimport libav as lib
+
+
+cdef class IndexEntries:
+    cdef lib.AVStream *stream_ptr
+    cdef _init(self, lib.AVStream *ptr)
+
+
+cdef IndexEntries wrap_index_entries(lib.AVStream *ptr)
diff --git a/av/indexentries.py b/av/indexentries.py
@@ -0,0 +1,91 @@
+import cython
+import cython.cimports.libav as lib
+from cython.cimports.av.indexentry import wrap_index_entry
+from cython.cimports.libc.stdint import int64_t
+
+_cinit_bypass_sentinel = cython.declare(object, object())
+
+
+@cython.cfunc
+def wrap_index_entries(ptr: cython.pointer[lib.AVStream]) -> IndexEntries:
+    obj: IndexEntries = IndexEntries(_cinit_bypass_sentinel)
+    obj._init(ptr)
+    return obj
+
+
+@cython.cclass
+class IndexEntries:
+    """A sequence-like view of FFmpeg's per-stream index entries.
+
+    Exposed as :attr:`~av.stream.Stream.index_entries`.
+
+    The index is provided by the demuxer and may be empty or incomplete depending
+    on the container format. This is useful for fast multi-seek loops (e.g., decoding
+    at a lower-than-native framerate).
+    """
+
+    def __cinit__(self, sentinel):
+        if sentinel is _cinit_bypass_sentinel:
+            return
+        raise RuntimeError("cannot manually instantiate IndexEntries")
+
+    @cython.cfunc
+    def _init(self, ptr: cython.pointer[lib.AVStream]):
+        self.stream_ptr = ptr
+
+    def __repr__(self):
+        return f"<av.IndexEntries[{len(self)}]>"
+
+    def __len__(self) -> int:
+        with cython.nogil:
+            return lib.avformat_index_get_entries_count(self.stream_ptr)
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield self[i]
+
+    def __getitem__(self, index):
+        if isinstance(index, int):
+            n = len(self)
+            if index < 0:
+                index += n
+            if index < 0 or index >= n:
+                raise IndexError(f"Index entries {index} out of bounds for size {n}")
+
+            c_idx = cython.declare(cython.int, index)
+            with cython.nogil:
+                entry = lib.avformat_index_get_entry(self.stream_ptr, c_idx)
+
+            if entry == cython.NULL:
+                raise IndexError("index entry not found")
+
+            return wrap_index_entry(entry)
+
+        elif isinstance(index, slice):
+            start, stop, step = index.indices(len(self))
+            return [self[i] for i in range(start, stop, step)]
+
+        else:
+            raise TypeError("Index must be an integer or a slice")
+
+    def search_timestamp(
+        self, timestamp, *, backward: bool = True, any_frame: bool = False
+    ):
+        """Search the underlying index for ``timestamp``.
+
+        This wraps FFmpeg's ``av_index_search_timestamp``.
+
+        Returns an index into this object, or ``-1`` if no match is found.
+        """
+        c_timestamp = cython.declare(int64_t, timestamp)
+        flags = cython.declare(cython.int, 0)
+
+        if backward:
+            flags |= lib.AVSEEK_FLAG_BACKWARD
+        if any_frame:
+            flags |= lib.AVSEEK_FLAG_ANY
+
+        with cython.nogil:
+            idx = lib.av_index_search_timestamp(self.stream_ptr, c_timestamp, flags)
+
+        return idx
diff --git a/av/indexentries.pyi b/av/indexentries.pyi
@@ -0,0 +1,14 @@
+from typing import Iterator, overload
+
+from av.indexentry import IndexEntry
+
+class IndexEntries:
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[IndexEntry]: ...
+    @overload
+    def __getitem__(self, index: int) -> IndexEntry: ...
+    @overload
+    def __getitem__(self, index: slice) -> list[IndexEntry]: ...
+    def search_timestamp(
+        self, timestamp, *, backward: bool = True, any_frame: bool = False
+    ) -> int: ...
diff --git a/av/indexentry.pxd b/av/indexentry.pxd
@@ -0,0 +1,9 @@
+cimport libav as lib
+
+
+cdef class IndexEntry:
+    cdef lib.AVIndexEntry *ptr
+    cdef _init(self, lib.AVIndexEntry *ptr)
+
+
+cdef IndexEntry wrap_index_entry(lib.AVIndexEntry *ptr)
diff --git a/av/indexentry.py b/av/indexentry.py
@@ -0,0 +1,63 @@
+import cython
+import cython.cimports.libav as lib
+
+_cinit_bypass_sentinel = cython.declare(object, object())
+
+
+@cython.cfunc
+def wrap_index_entry(ptr: cython.pointer[lib.AVIndexEntry]) -> IndexEntry:
+    obj: IndexEntry = IndexEntry(_cinit_bypass_sentinel)
+    obj._init(ptr)
+    return obj
+
+
+@cython.cclass
+class IndexEntry:
+    """A single entry from a stream's index.
+
+    This is a thin wrapper around FFmpeg's ``AVIndexEntry``.
+
+    The exact meaning of the fields depends on the container/demuxer.
+    """
+
+    def __cinit__(self, sentinel):
+        if sentinel is not _cinit_bypass_sentinel:
+            raise RuntimeError("cannot manually instantiate IndexEntry")
+
+    @cython.cfunc
+    def _init(self, ptr: cython.pointer[lib.AVIndexEntry]):
+        self.ptr = ptr
+
+    def __repr__(self):
+        return (
+            f"<av.IndexEntry pos={self.pos} timestamp={self.timestamp} flags={self.flags} "
+            f"size={self.size} min_distance={self.min_distance}>"
+        )
+
+    @property
+    def pos(self):
+        return self.ptr.pos
+
+    @property
+    def timestamp(self):
+        return self.ptr.timestamp
+
+    @property
+    def flags(self):
+        return self.ptr.flags
+
+    @property
+    def is_keyframe(self):
+        return bool(self.ptr.flags & lib.AVINDEX_KEYFRAME)
+
+    @property
+    def is_discard(self):
+        return bool(self.ptr.flags & lib.AVINDEX_DISCARD_FRAME)
+
+    @property
+    def size(self):
+        return self.ptr.size
+
+    @property
+    def min_distance(self):
+        return self.ptr.min_distance
diff --git a/av/indexentry.pyi b/av/indexentry.pyi
@@ -0,0 +1,8 @@
+class IndexEntry:
+    pos: int
+    timestamp: int
+    flags: int
+    is_keyframe: bool
+    is_discard: bool
+    size: int
+    min_distance: int
diff --git a/av/stream.pxd b/av/stream.pxd
@@ -3,6 +3,7 @@ cimport libav as lib
 from av.codec.context cimport CodecContext
 from av.container.core cimport Container
 from av.frame cimport Frame
+from av.indexentries cimport IndexEntries
 from av.packet cimport Packet
 
 
@@ -16,6 +17,8 @@ cdef class Stream:
     # CodecContext attributes.
     cdef readonly CodecContext codec_context
 
+    cdef readonly IndexEntries index_entries
+
     # Private API.
     cdef _init(self, Container, lib.AVStream*, CodecContext)
     cdef _finalize_for_output(self)
diff --git a/av/stream.py b/av/stream.py
@@ -3,6 +3,7 @@
 import cython
 from cython.cimports import libav as lib
 from cython.cimports.av.error import err_check
+from cython.cimports.av.indexentries import wrap_index_entries
 from cython.cimports.av.packet import Packet
 from cython.cimports.av.utils import (
     avdict_to_dict,
@@ -106,6 +107,7 @@ def _init(
     ):
         self.container = container
         self.ptr = stream
+        self.index_entries = wrap_index_entries(self.ptr)
 
         self.codec_context = codec_context
         if self.codec_context:
diff --git a/av/stream.pyi b/av/stream.pyi
@@ -4,6 +4,7 @@ from typing import Literal, cast
 
 from .codec import Codec, CodecContext
 from .container import Container
+from .indexentries import IndexEntries
 
 class Disposition(Flag):
     default = cast(int, ...)
@@ -32,6 +33,7 @@ class Stream:
     codec: Codec
     codec_context: CodecContext
     metadata: dict[str, str]
+    index_entries: IndexEntries
     id: int
     profiles: list[str]
     profile: str | None
diff --git a/include/libavformat/avformat.pxd b/include/libavformat/avformat.pxd
@@ -331,3 +331,18 @@ cdef extern from "libavformat/avformat.h" nogil:
     # custom
 
     cdef set pyav_get_available_formats()
+
+    cdef struct AVIndexEntry:
+        int64_t pos
+        int64_t timestamp
+        int flags
+        int size
+        int min_distance
+
+    cdef enum:
+        AVINDEX_KEYFRAME
+        AVINDEX_DISCARD_FRAME
+
+    cdef AVIndexEntry *avformat_index_get_entry(AVStream *st, int idx)
+    cdef int avformat_index_get_entries_count(AVStream *st)
+    cdef int av_index_search_timestamp(AVStream *st, int64_t timestamp, int flags)
diff --git a/tests/test_indexentries.py b/tests/test_indexentries.py
@@ -0,0 +1,79 @@
+import av
+
+from .common import TestCase, fate_suite
+
+
+class TestIndexEntries(TestCase):
+    def test_index_entries_len_mp4(self) -> None:
+        with av.open(fate_suite("h264/interlaced_crop.mp4")) as container:
+            stream = container.streams.video[0]
+            assert len(stream.index_entries) == stream.frames
+
+    def test_index_entries_len_webm(self) -> None:
+        with av.open(
+            fate_suite("vp9-test-vectors/vp90-2-00-quantizer-00.webm")
+        ) as container:
+            stream = container.streams.video[0]
+            index_entries_len_before_demux = len(stream.index_entries)
+
+            keyframes = len([p for p in container.demux(video=0) if p.is_keyframe])
+            assert index_entries_len_before_demux == keyframes
+
+    def test_index_entries_search_timestamp_options_mp4(self) -> None:
+        with av.open(fate_suite("h264/interlaced_crop.mp4")) as container:
+            stream = container.streams.video[0]
+            fi = stream.index_entries
+
+            target_ts = -1
+
+            assert fi.search_timestamp(target_ts) == 0
+            assert fi.search_timestamp(target_ts, any_frame=True) == 1
+            assert fi.search_timestamp(target_ts, backward=False) == 21
+            assert fi.search_timestamp(target_ts, backward=False, any_frame=True) == 1
+
+            e0 = fi[0]
+            e1 = fi[1]
+            e21 = fi[21]
+            assert e0.timestamp == -2 and e0.is_keyframe
+            assert e1.timestamp == -1 and not e1.is_keyframe
+            assert e21.timestamp == 19 and e21.is_keyframe
+
+    def test_index_entries_matches_packet_mp4(self) -> None:
+        with av.open(fate_suite("h264/interlaced_crop.mp4")) as container:
+            stream = container.streams.video[0]
+            for i, packet in enumerate(container.demux(video=0)):
+                if packet.dts is not None:
+                    assert stream.index_entries[i].timestamp == packet.dts
+
+    def test_index_entries_in_bounds(self) -> None:
+        with av.open(fate_suite("h264/interlaced_crop.mp4")) as container:
+            stream = container.streams.video[0]
+            first = stream.index_entries[0]
+            first_neg = stream.index_entries[-len(stream.index_entries)]
+            last = stream.index_entries[-1]
+            last_pos = stream.index_entries[len(stream.index_entries) - 1]
+            assert first.timestamp == first_neg.timestamp
+            assert last.timestamp == last_pos.timestamp
+
+    def test_index_entries_out_of_bounds(self) -> None:
+        with av.open(fate_suite("h264/interlaced_crop.mp4")) as container:
+            stream = container.streams.video[0]
+            with self.assertRaises(IndexError):
+                _ = stream.index_entries[len(stream.index_entries)]
+
+            with self.assertRaises(IndexError):
+                _ = stream.index_entries[-len(stream.index_entries) - 1]
+
+    def test_index_entries_slice(self) -> None:
+        with av.open(fate_suite("h264/interlaced_crop.mp4")) as container:
+            stream = container.streams.video[0]
+
+            individual_indices = [stream.index_entries[i] for i in range(1, 5)]
+            slice_indices = stream.index_entries[1:5]
+            assert len(individual_indices) == len(slice_indices) == 4
+            assert all(
+                [
+                    i.timestamp == j.timestamp
+                    for i, j in zip(individual_indices, slice_indices)
+                ]
+            )
