diff --git a/av/audio/__init__.pyi b/av/audio/__init__.pyi index 73f2eebdd..daefab6c9 100644 --- a/av/audio/__init__.pyi +++ b/av/audio/__init__.pyi @@ -1,4 +1,16 @@ +from typing import Literal + from .frame import AudioFrame from .stream import AudioStream +_AudioCodecName = Literal[ + "aac", + "libopus", + "mp2", + "mp3", + "pcm_alaw", + "pcm_mulaw", + "pcm_s16le", +] + __all__ = ("AudioFrame", "AudioStream") diff --git a/av/codec/context.pyi b/av/codec/context.pyi index ac57126ef..77810d9ed 100644 --- a/av/codec/context.pyi +++ b/av/codec/context.pyi @@ -1,8 +1,12 @@ from enum import Flag, IntEnum from fractions import Fraction -from typing import ClassVar, Literal, cast +from typing import ClassVar, Literal, cast, overload +from av.audio import _AudioCodecName +from av.audio.codeccontext import AudioCodecContext from av.packet import Packet +from av.video import _VideoCodecName +from av.video.codeccontext import VideoCodecContext from .codec import Codec from .hwaccel import HWAccel @@ -87,6 +91,21 @@ class CodecContext: @property def is_hwaccel(self) -> bool: ... def open(self, strict: bool = True) -> None: ... + @overload + @staticmethod + def create( + codec: _AudioCodecName, + mode: Literal["r", "w"] | None = None, + hwaccel: HWAccel | None = None, + ) -> AudioCodecContext: ... + @overload + @staticmethod + def create( + codec: _VideoCodecName, + mode: Literal["r", "w"] | None = None, + hwaccel: HWAccel | None = None, + ) -> VideoCodecContext: ... + @overload @staticmethod def create( codec: str | Codec, diff --git a/av/container/output.pyi b/av/container/output.pyi index 3169f486b..568345cd2 100644 --- a/av/container/output.pyi +++ b/av/container/output.pyi @@ -1,11 +1,12 @@ from fractions import Fraction -from typing import Literal, Sequence, TypeVar, Union, overload +from typing import Sequence, TypeVar, Union, overload +from av.audio import _AudioCodecName from av.audio.stream import AudioStream from av.data.stream import DataStream from av.packet import Packet -from av.stream import Stream from av.subtitles.stream import SubtitleStream +from av.video import _VideoCodecName from av.video.stream import VideoStream from .core import Container @@ -17,7 +18,7 @@ class OutputContainer(Container): @overload def add_stream( self, - codec_name: Literal["pcm_s16le", "aac", "mp3", "mp2"], + codec_name: _AudioCodecName, rate: int | None = None, options: dict[str, str] | None = None, **kwargs, @@ -25,7 +26,7 @@ class OutputContainer(Container): @overload def add_stream( self, - codec_name: Literal["h264", "hevc", "mpeg4", "png", "gif", "qtrle"], + codec_name: _VideoCodecName, rate: Fraction | int | None = None, options: dict[str, str] | None = None, **kwargs, diff --git a/av/video/__init__.pyi b/av/video/__init__.pyi index 8fa8fe7e5..58a19a63f 100644 --- a/av/video/__init__.pyi +++ b/av/video/__init__.pyi @@ -1,4 +1,17 @@ +from typing import Literal + from .frame import VideoFrame from .stream import VideoStream +_VideoCodecName = Literal[ + "gif", + "h264", + "hevc", + "libvpx", + "libx264", + "mpeg4", + "png", + "qtrle", +] + __all__ = ("VideoFrame", "VideoStream")