Add Graph.vpush() and Graph.vpull()

Author: WyattBlue

av/container/input.pyi

@@ -21,6 +21,12 @@ class InputContainer(Container):
     def close(self) -> None: ...
     def demux(self, *args: Any, **kwargs: Any) -> Iterator[Packet]: ...
     @overload
+    def decode(self, video: int) -> Iterator[VideoFrame]: ...
+    @overload
+    def decode(self, audio: int) -> Iterator[AudioFrame]: ...
+    @overload
+    def decode(self, subtitles: int) -> Iterator[SubtitleSet]: ...
+    @overload
     def decode(self, *args: VideoStream) -> Iterator[VideoFrame]: ...
     @overload
     def decode(self, *args: AudioStream) -> Iterator[AudioFrame]: ...

av/filter/graph.pxd

@@ -4,7 +4,6 @@ from av.filter.context cimport FilterContext
 
 
 cdef class Graph:
-
     cdef lib.AVFilterGraph *ptr
 
     cdef readonly bint configured

av/filter/graph.pyi

@@ -41,3 +41,5 @@ class Graph:
     ) -> FilterContext: ...
     def push(self, frame: None | AudioFrame | VideoFrame) -> None: ...
     def pull(self) -> VideoFrame | AudioFrame: ...
+    def vpush(self, frame: VideoFrame | None) -> None: ...
+    def vpull(self) -> VideoFrame: ...

av/filter/graph.pyx

@@ -174,11 +174,15 @@ cdef class Graph:
         else:
             raise ValueError(f"can only AudioFrame, VideoFrame or None; got {type(frame)}")
 
-        if len(contexts) != 1:
-            raise ValueError(f"can only auto-push with single buffer; found {len(contexts)}")
+        for ctx in contexts:
+            ctx.push(frame)
 
-        contexts[0].push(frame)
+    def vpush(self, VideoFrame frame):
+        for ctx in self._context_by_type.get("buffer", []):
+            ctx.push(frame)
 
+
+    # TODO: Test complex filter graphs, add `at: int = 0` arg to pull() and vpull().
     def pull(self):
         vsinks = self._context_by_type.get("buffersink", [])
         asinks = self._context_by_type.get("abuffersink", [])
@@ -188,3 +192,11 @@ cdef class Graph:
             raise ValueError(f"can only auto-pull with single sink; found {nsinks}")
 
         return (vsinks or asinks)[0].pull()
+
+    def vpull(self):
+        vsinks = self._context_by_type.get("buffersink", [])
+        nsinks = len(vsinks)
+        if nsinks != 1:
+            raise ValueError(f"can only auto-pull with single sink; found {nsinks}")
+
+        return vsinks[0].pull()

tests/test_filters.py

@@ -204,7 +204,7 @@ def test_video_buffer(self):
 
         for frame in input_container.decode():
             self.assertEqual(frame.time_base, Fraction(1, 30))
-            graph.push(frame)
+            graph.vpush(frame)
             filtered_frames = pull_until_blocked(graph)
 
             if frame.pts == 0:
@@ -220,7 +220,7 @@ def test_video_buffer(self):
                 self.assertEqual(filtered_frames[1].pts, (frame.pts - 1) * 2 + 1)
                 self.assertEqual(filtered_frames[1].time_base, Fraction(1, 60))
 
-    def test_EOF(self):
+    def test_EOF(self) -> None:
         input_container = av.open(format="lavfi", file="color=c=pink:duration=1:r=30")
         video_stream = input_container.streams.video[0]
 
@@ -233,12 +233,12 @@ def test_EOF(self):
         graph.configure()
 
         for frame in input_container.decode(video=0):
-            graph.push(frame)
+            graph.vpush(frame)
 
-        graph.push(None)
+        graph.vpush(None)
 
         # if we do not push None, we get a BlockingIOError
-        palette_frame = graph.pull()
+        palette_frame = graph.vpull()
 
         self.assertIsInstance(palette_frame, av.VideoFrame)
         self.assertEqual(palette_frame.width, 16)