Improve InputContainer.decode type stubs

av/container/input.pyi

@@ -1,8 +1,10 @@
 from typing import Iterator, Literal, overload
 
+from av.audio.frame import AudioFrame
 from av.audio.stream import AudioStream
-from av.frame import Frame
 from av.packet import Packet
+from av.subtitles.stream import SubtitleStream
+from av.subtitles.subtitle import SubtitleSet
 from av.video.frame import VideoFrame
 from av.video.stream import VideoStream
 
@@ -13,12 +15,19 @@ class InputContainer(Container):
     bit_rate: int
     size: int
 
+    def __enter__(self) -> InputContainer: ...
     def close(self) -> None: ...
     def demux(self, *args, **kwargs) -> Iterator[Packet]: ...
     @overload
-    def decode(self, *args: VideoStream, **kwargs) -> Iterator[VideoFrame]: ...
+    def decode(self, *args: VideoStream) -> Iterator[VideoFrame]: ...
     @overload
-    def decode(self, *args, **kwargs) -> Iterator[Frame]: ...
+    def decode(self, *args: AudioStream) -> Iterator[AudioFrame]: ...
+    @overload
+    def decode(self, *args: SubtitleStream) -> Iterator[SubtitleSet]: ...
+    @overload
+    def decode(
+        self, *args, **kwargs
+    ) -> Iterator[VideoFrame | AudioFrame | SubtitleSet]: ...
     def seek(
         self,
         offset: int,

av/container/output.pyi

@@ -1,4 +1,5 @@
 from .core import Container
 
 class OutputContainer(Container):
+    def __enter__(self) -> OutputContainer: ...
     def start_encoding(self) -> None: ...