From 2b6586c3e49587a44cb56d52380be8bb58c92e79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jeremy=20Lain=C3=A9?= <jeremy.laine@m4x.org>
Date: Tue, 1 Apr 2025 09:28:04 +0200
Subject: [PATCH] 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
---
 av/audio/__init__.pyi   | 12 ++++++++++++
 av/codec/context.pyi    | 21 ++++++++++++++++++++-
 av/container/output.pyi |  9 +++++----
 av/video/__init__.pyi   | 13 +++++++++++++
 4 files changed, 50 insertions(+), 5 deletions(-)

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