Improve stubs in av.container

av/container/core.pyi

@@ -1,11 +1,31 @@
 from numbers import Real
 from pathlib import Path
-from typing import Any, Iterator, Literal, overload
+from types import TracebackType
+from typing import Any, Callable, Literal, Type, overload
+
+from av.enum import EnumFlag
 
 from .input import InputContainer
 from .output import OutputContainer
 from .streams import StreamContainer
 
+class Flags(EnumFlag):
+    GENPTS: int
+    IGNIDX: int
+    NONBLOCK: int
+    IGNDTS: int
+    NOFILLIN: int
+    NOPARSE: int
+    NOBUFFER: int
+    CUSTOM_IO: int
+    DISCARD_CORRUPT: int
+    FLUSH_PACKETS: int
+    BITEXACT: int
+    SORT_DTS: int
+    FAST_SEEK: int
+    SHORTEST: int
+    AUTO_BSF: int
+
 class Container:
     writeable: bool
     name: str
@@ -21,13 +41,17 @@ class Container:
     container_options: dict[str, str]
     stream_options: list[str]
     streams: StreamContainer
-    duration: int | None
     metadata: dict[str, str]
     open_timeout: Real | None
     read_timeout: Real | None
 
     def __enter__(self) -> Container: ...
-    def __exit__(self, exc_type, exc_val, exc_tb): ...
+    def __exit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool: ...
     def err_check(self, value: int) -> int: ...
     def set_timeout(self, timeout: Real | None) -> None: ...
     def start_timeout(self) -> None: ...
@@ -43,8 +67,8 @@ def open(
     metadata_encoding: str = "utf-8",
     metadata_errors: str = "strict",
     buffer_size: int = 32768,
-    timeout=Real | None | tuple[Real | None, Real | None],
-    io_open=None,
+    timeout: Real | None | tuple[Real | None, Real | None] = None,
+    io_open: Callable[..., Any] | None = None,
 ) -> InputContainer: ...
 @overload
 def open(
@@ -57,8 +81,8 @@ def open(
     metadata_encoding: str = "utf-8",
     metadata_errors: str = "strict",
     buffer_size: int = 32768,
-    timeout=Real | None | tuple[Real | None, Real | None],
-    io_open=None,
+    timeout: Real | None | tuple[Real | None, Real | None] = None,
+    io_open: Callable[..., Any] | None = None,
 ) -> InputContainer: ...
 @overload
 def open(
@@ -71,8 +95,8 @@ def open(
     metadata_encoding: str = "utf-8",
     metadata_errors: str = "strict",
     buffer_size: int = 32768,
-    timeout=Real | None | tuple[Real | None, Real | None],
-    io_open=None,
+    timeout: Real | None | tuple[Real | None, Real | None] = None,
+    io_open: Callable[..., Any] | None = None,
 ) -> OutputContainer: ...
 @overload
 def open(
@@ -85,6 +109,6 @@ def open(
     metadata_encoding: str = "utf-8",
     metadata_errors: str = "strict",
     buffer_size: int = 32768,
-    timeout=Real | None | tuple[Real | None, Real | None],
-    io_open=None,
+    timeout: Real | None | tuple[Real | None, Real | None] = None,
+    io_open: Callable[..., Any] | None = None,
 ) -> InputContainer | OutputContainer: ...

av/container/input.pyi

@@ -1,23 +1,25 @@
-from typing import Iterator, Literal, overload
+from typing import Any, Iterator, Literal, overload
 
 from av.audio.frame import AudioFrame
 from av.audio.stream import AudioStream
 from av.packet import Packet
+from av.stream import Stream
 from av.subtitles.stream import SubtitleStream
 from av.subtitles.subtitle import SubtitleSet
 from av.video.frame import VideoFrame
 from av.video.stream import VideoStream
 
 from .core import Container
-from .streams import Stream
 
 class InputContainer(Container):
+    start_time: int
+    duration: int | None
     bit_rate: int
     size: int
 
     def __enter__(self) -> InputContainer: ...
     def close(self) -> None: ...
-    def demux(self, *args, **kwargs) -> Iterator[Packet]: ...
+    def demux(self, *args: Any, **kwargs: Any) -> Iterator[Packet]: ...
     @overload
     def decode(self, *args: VideoStream) -> Iterator[VideoFrame]: ...
     @overload
@@ -26,7 +28,7 @@ class InputContainer(Container):
     def decode(self, *args: SubtitleStream) -> Iterator[SubtitleSet]: ...
     @overload
     def decode(
-        self, *args, **kwargs
+        self, *args: Any, **kwargs: Any
     ) -> Iterator[VideoFrame | AudioFrame | SubtitleSet]: ...
     def seek(
         self,
@@ -39,3 +41,4 @@ class InputContainer(Container):
         unsupported_frame_offset: bool = False,
         unsupported_byte_offset: bool = False,
     ) -> None: ...
+    def flush_buffers(self) -> None: ...

av/container/output.pyi

@@ -1,5 +1,21 @@
+from fractions import Fraction
+from typing import Sequence
+
+from av.packet import Packet
+from av.stream import Stream
+
 from .core import Container
 
 class OutputContainer(Container):
     def __enter__(self) -> OutputContainer: ...
+    def add_stream(
+        self,
+        codec_name: str | None = None,
+        rate: Fraction | int | float | None = None,
+        template: Stream | None = None,
+        options: dict[str, str] | None = None,
+    ) -> Stream: ...
     def start_encoding(self) -> None: ...
+    def close(self) -> None: ...
+    def mux(self, packets: Sequence[Packet]) -> None: ...
+    def mux_one(self, packet: Packet) -> None: ...

av/container/streams.pyi

@@ -1,69 +1,16 @@
-from fractions import Fraction
-from typing import Iterator, Literal, overload
+from typing import Iterator, overload
 
+from av.audio.stream import AudioStream
+from av.data.stream import DataStream
+from av.stream import Stream
+from av.subtitles.stream import SubtitleStream
 from av.video.stream import VideoStream
 
-class Codec:
-    name: str
-    mode: Literal["r", "w"]
-
-    frame_rates: list[Fraction] | None
-    audio_rates: list[int] | None
-
-class CodecContext:
-    name: str
-    bit_rate: int | None
-    width: int
-    height: int
-    pix_fmt: str | None
-    sample_aspect_ratio: Fraction | None
-    sample_rate: int | None
-    channels: int
-    extradata_size: int
-    is_open: Literal[0, 1]
-    is_encoder: Literal[0, 1]
-    is_decoder: Literal[0, 1]
-
-class Stream:
-    thread_type: Literal["NONE", "FRAME", "SLICE", "AUTO"]
-
-    id: int
-    profile: str | None
-    codec_context: CodecContext
-
-    index: int
-    time_base: Fraction | None
-    average_rate: Fraction | None
-    base_rate: Fraction | None
-    guessed_rate: Fraction | None
-
-    start_time: int | None
-    duration: int | None
-    frames: int
-    language: str | None
-
-    # Defined by `av_get_media_type_string` at
-    # https://ffmpeg.org/doxygen/6.0/libavutil_2utils_8c_source.html
-    type: Literal["video", "audio", "data", "subtitle", "attachment"]
-
-    # From `codec_context`
-    name: str
-    bit_rate: int | None
-    sample_rate: int | None
-    channels: int
-    extradata_size: int
-    is_open: Literal[0, 1]
-    is_encoder: Literal[0, 1]
-    is_decoder: Literal[0, 1]
-
-    def decode(self, packet=None): ...
-    def encode(self, frame=None): ...
-
 class StreamContainer:
     video: tuple[VideoStream, ...]
-    audio: tuple[Stream, ...]
-    subtitles: tuple[Stream, ...]
-    data: tuple[Stream, ...]
+    audio: tuple[AudioStream, ...]
+    subtitles: tuple[SubtitleStream, ...]
+    data: tuple[DataStream, ...]
     other: tuple[Stream, ...]
 
     def __init__(self) -> None: ...