Fix more type stub mismatches

CHANGELOG.rst

@@ -38,6 +38,7 @@ Fixes:
 - ``CodecContext.is_encoder``, ``CodecContext.is_decoder``, ``Container.input_was_opened``, and ``VideoFormatComponent.is_chroma`` now return ``bool`` instead of ``int``.
 - ``Frame.time_base`` returns ``AVRational(0, 1)`` instead of ``None`` when unset, like other ``AVRational`` properties.
 - Fix type stubs that did not match runtime: ``average_rate``, ``base_rate``, and ``guessed_rate`` are only on ``VideoStream``; ``codec``, ``codec_tag``, and ``options`` are only on audio, video, and subtitle streams; ``Stream.codec_context`` may be ``None``; ``AudioStream.bit_rate``, ``InputContainer.start_time``, ``VideoFrame.time``, and ``FilterPad.context`` may be ``None``; ``Codec.hardware_configs`` is a tuple; ``CodecContext.delay`` is an ``int``; ``AudioStream.max_bit_rate`` is added; and the inaccessible ``Plane.frame``, ``Plane.index``, and ``Container.buffer_size`` are removed.
+- Fix more type stubs that did not match runtime: ``SubtitleSet`` items are ``AssSubtitle | BitmapSubtitle``; ``BitmapSubtitle`` supports ``len()``, iteration, and indexing; ``BitmapSubtitlePlane`` supports the buffer protocol; ``VideoFormat`` supports ``int()``; ``CodecContext.encode``, ``encode_lazy``, and ``decode`` have return types; ``SubtitleCodecContext.decode`` requires a packet; ``VideoStream.format`` may be ``None``; stream ``name`` is a ``str`` on audio, video, and subtitle streams; and ``VideoFrame.to_image``, ``set_image``, and ``from_image`` use ``PIL.Image.Image``.
 
 v19.0.2
 -------

av/audio/stream.pyi

@@ -26,6 +26,7 @@ class AudioStream(Stream):
     def decode(self, packet: Packet | None = None) -> list[AudioFrame]: ...
 
     # From codec context
+    name: str
     codec: Codec
     codec_tag: str
     options: dict[str, str]

av/codec/context.pyi

@@ -1,3 +1,4 @@
+from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from enum import Flag, IntEnum, IntFlag
 from fractions import Fraction
@@ -13,6 +14,7 @@ from av.rational import AVRational
 from av.sidedata.sidedata import Type as FrameSideDataType
 from av.subtitles import _SubtitleCodecName
 from av.subtitles.codeccontext import SubtitleCodecContext
+from av.subtitles.subtitle import AssSubtitle, BitmapSubtitle
 from av.video import _VideoCodecName
 from av.video.codeccontext import VideoCodecContext
 
@@ -227,6 +229,8 @@ class CodecContext:
         self, raw_input: bytes | bytearray | memoryview | None = None
     ) -> list[Packet]: ...
     def flush_buffers(self) -> None: ...
-    def encode(self, frame: Frame | None = None): ...
-    def encode_lazy(self, frame: Frame | None = None): ...
-    def decode(self, packet: Packet | None = None): ...
+    def encode(self, frame: Frame | None = None) -> list[Packet]: ...
+    def encode_lazy(self, frame: Frame | None = None) -> Iterator[Packet]: ...
+    def decode(
+        self, packet: Packet | None = None
+    ) -> Sequence[Frame | AssSubtitle | BitmapSubtitle]: ...

av/subtitles/codeccontext.pyi

@@ -2,11 +2,14 @@ from typing import Literal, final
 
 from av.codec.context import CodecContext
 from av.packet import Packet
-from av.subtitles.subtitle import SubtitleSet
+from av.subtitles.subtitle import AssSubtitle, BitmapSubtitle, SubtitleSet
 
 @final
 class SubtitleCodecContext(CodecContext):
     type: Literal["subtitle"]
     subtitle_header: bytes | None
+    def decode(  # type: ignore[override]
+        self, packet: Packet
+    ) -> list[AssSubtitle] | list[BitmapSubtitle]: ...
     def decode2(self, packet: Packet) -> SubtitleSet | None: ...
     def encode_subtitle(self, subtitle: SubtitleSet) -> Packet: ...

av/subtitles/stream.pyi

@@ -12,6 +12,7 @@ class SubtitleStream(Stream):
     codec_context: SubtitleCodecContext
 
     # From codec context
+    name: str
     codec: Codec
     codec_tag: str
     options: dict[str, str]

av/subtitles/subtitle.pyi

@@ -12,7 +12,7 @@ class SubtitleSet:
     start_display_time: int
     end_display_time: int
     pts: int
-    rects: tuple[Subtitle, ...]
+    rects: tuple[AssSubtitle | BitmapSubtitle, ...]
 
     @staticmethod
     def create(
@@ -23,8 +23,8 @@ class SubtitleSet:
         subtitle_format: int = 1,
     ) -> SubtitleSet: ...
     def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[Subtitle]: ...
-    def __getitem__(self, i: int, /) -> Subtitle: ...
+    def __iter__(self) -> Iterator[AssSubtitle | BitmapSubtitle]: ...
+    def __getitem__(self, i: int, /) -> AssSubtitle | BitmapSubtitle: ...
 
 @disjoint_base
 class Subtitle:
@@ -39,12 +39,16 @@ class BitmapSubtitle(Subtitle):
     height: int
     nb_colors: int
     planes: tuple[BitmapSubtitlePlane, ...]
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[BitmapSubtitlePlane]: ...
+    def __getitem__(self, i: int, /) -> BitmapSubtitlePlane: ...
 
 @final
 class BitmapSubtitlePlane:
     subtitle: BitmapSubtitle
     index: int
     buffer_size: int
+    def __buffer__(self, flags: int, /) -> memoryview: ...
 
 @final
 class AssSubtitle(Subtitle):

av/video/format.pyi

@@ -20,6 +20,7 @@ class VideoFormat:
     def __init__(self, name: str, width: int = 0, height: int = 0) -> None: ...
     def chroma_width(self, luma_width: int = 0) -> int: ...
     def chroma_height(self, luma_height: int = 0) -> int: ...
+    def __int__(self) -> int: ...
 
 @final
 class VideoFormatComponent:

av/video/frame.pyi

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import Any, final
 
 import numpy as np
+from PIL.Image import Image
 
 from av.frame import Frame
 
@@ -89,13 +90,13 @@ class VideoFrame(Frame):
     ) -> VideoFrame: ...
     def to_rgb(self, **kwargs: Any) -> VideoFrame: ...
     def save(self, filepath: str | Path, **options: Any) -> None: ...
-    def set_image(self, img: Any) -> None: ...
-    def to_image(self, **kwargs): ...
+    def set_image(self, img: Image) -> None: ...
+    def to_image(self, **kwargs: Any) -> Image: ...
     def to_ndarray(
         self, channel_last: bool = False, **kwargs: Any
     ) -> _SupportedNDarray: ...
     @staticmethod
-    def from_image(img): ...
+    def from_image(img: Image) -> VideoFrame: ...
     @staticmethod
     def from_numpy_buffer(
         array: _SupportedNDarray, format: str = "rgb24", width: int = 0

av/video/stream.pyi

@@ -33,10 +33,11 @@ class VideoStream(Stream):
     guessed_rate: AVRational
 
     # from codec context
+    name: str
     codec: Codec
     codec_tag: str
     options: dict[str, str]
-    format: VideoFormat
+    format: VideoFormat | None
     thread_count: int
     @property
     def thread_type(self) -> ThreadType: ...

examples/subtitles/remux.py

@@ -1,4 +1,5 @@
 import av
+from av.subtitles.subtitle import AssSubtitle
 
 av.logging.set_level(av.logging.VERBOSE)
 
@@ -22,4 +23,5 @@
 with av.open("remuxed.vtt") as f:
     for subset in f.decode(subtitles=0):
         for sub in subset:
-            print(sub.ass)
+            if isinstance(sub, AssSubtitle):
+                print(sub.ass)

tests/test_encode.py

@@ -110,6 +110,7 @@ def assert_rgb_rotate(
     # codec context properties
     assert stream.codec.name == "mpeg4"
     assert stream.codec.long_name == "MPEG-4 part 2"
+    assert stream.format is not None
     assert stream.format.name == "yuv420p"
     assert stream.format.width == WIDTH
     assert stream.format.height == HEIGHT
@@ -124,6 +125,7 @@ def test_default_options(self) -> None:
             assert not stream.time_base
 
             # codec context properties
+            assert stream.format is not None
             assert stream.format.height == 480
             assert stream.format.name == "yuv420p"
             assert stream.format.width == 640

tests/test_file_probing.py

@@ -249,6 +249,7 @@ def test_stream_probing(self) -> None:
         assert stream.codec.long_name == "MPEG-2 video"
         assert stream.codec.name == "mpeg2video"
         assert stream.display_aspect_ratio == Fraction(4, 3)
+        assert stream.format is not None
         assert stream.format.name == "yuv420p"
         assert not stream.has_b_frames
         assert stream.height == 576

tests/test_python_io.py

@@ -293,6 +293,7 @@ def test_writing_to_custom_io_image2(self) -> None:
 
             # codec context properties
             assert stream.codec.name == "png"
+            assert stream.format is not None
             assert stream.format.name == "rgb24"
             assert stream.format.width == width
             assert stream.format.height == height