Add typing overloads for CodecContext.create

The return type of `CodecContext.create` can be narrowed based on the
name of the requested codec, allowing us to know whether we are dealing
with an `AudioCodecContext` or a `VideoCodecContext`.

As the same logic is used by `OutputContainer.add_stream`, use a shared
list of audio / video codec names. We add the following codec names:

- audio : libopus, pcm_alaw, pcm_mulaw
- video : libvpx, libx264

Author: jlaine

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")

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,

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,15 +18,15 @@ 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,
     ) -> AudioStream: ...
     @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,

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")