Improve stub in av.audio

av/audio/codeccontext.pyi

@@ -1,11 +1,8 @@
 from typing import Literal
 
 from av.codec.context import CodecContext
-from av.frame import Frame
-from av.packet import Packet
 
 from .format import AudioFormat
-from .frame import AudioFrame
 from .layout import AudioLayout
 
 class AudioCodecContext(CodecContext):

av/audio/format.pyi

@@ -1,11 +1,9 @@
-from typing import Any
-
 class AudioFormat:
     name: str
     bytes: int
     bits: int
     is_planar: bool
     is_packed: bool
-    planar: Any
-    packed: Any
+    planar: AudioFormat
+    packed: AudioFormat
     container_name: str

av/audio/frame.pyi

@@ -1,8 +1,19 @@
+from typing import Any, Union
+
+import numpy as np
+
 from av.frame import Frame
 
 from .plane import AudioPlane
 
 format_dtypes: dict[str, str]
+_SupportedNDarray = Union[
+    np.ndarray[Any, np.dtype[np.float64]],  # f8
+    np.ndarray[Any, np.dtype[np.float32]],  # f4
+    np.ndarray[Any, np.dtype[np.int32]],  # i4
+    np.ndarray[Any, np.dtype[np.int16]],  # i2
+    np.ndarray[Any, np.dtype[np.uint8]],  # u1
+]
 
 class AudioFrame(Frame):
     planes: tuple[AudioPlane, ...]
@@ -16,9 +27,9 @@ class AudioFrame(Frame):
         layout: str = "stereo",
         samples: int = 0,
         align: int = 1,
-    ): ...
-    def to_ndarray(self): ...
+    ) -> None: ...
     @staticmethod
     def from_ndarray(
-        array, format: str = "s16", layout: str = "stereo"
+        array: _SupportedNDarray, format: str = "s16", layout: str = "stereo"
     ) -> AudioFrame: ...
+    def to_ndarray(self) -> _SupportedNDarray: ...

av/audio/layout.pyi

@@ -1,10 +1,11 @@
 channel_descriptions: dict[str, str]
 
 class AudioLayout:
+    name: str
     layout: int
     nb_channels: int
-    channels: tuple[AudioChannels]
+    channels: tuple[AudioChannel, ...]
 
-class AudioChannels:
+class AudioChannel:
     name: str
     description: str

av/audio/plane.pyi

@@ -1,2 +1,4 @@
-class AudioPlane:
+from av.plane import Plane
+
+class AudioPlane(Plane):
     buffer_size: int

av/audio/resampler.pyi

@@ -1,17 +1,20 @@
+from av.filter.graph import Graph
+
 from .format import AudioFormat
 from .frame import AudioFrame
+from .layout import AudioLayout
 
 class AudioResampler:
     rate: int
     frame_size: int
     format: AudioFormat
-    graph: None
+    graph: Graph | None
 
     def __init__(
         self,
-        format=None,
-        layout=None,
+        format: AudioFormat | None = None,
+        layout: AudioLayout | None = None,
         rate: int | None = None,
         frame_size: int | None = None,
-    ): ...
-    def resample(self, frame: AudioFrame | None) -> list: ...
+    ) -> None: ...
+    def resample(self, frame: AudioFrame | None) -> list[AudioFrame]: ...

av/audio/stream.pyi

@@ -1,3 +1,12 @@
+from av.packet import Packet
 from av.stream import Stream
 
-class AudioStream(Stream): ...
+from .format import AudioFormat
+from .frame import AudioFrame
+
+class AudioStream(Stream):
+    format: AudioFormat
+    type = "audio"
+
+    def encode(self, frame: AudioFrame | None = None) -> list[Packet]: ...
+    def decode(self, packet: Packet | None = None) -> list[AudioFrame]: ...