diff --git a/av/container/input.pyi b/av/container/input.pyi index f85eaec46..90154c331 100644 --- a/av/container/input.pyi +++ b/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]: ... diff --git a/av/filter/graph.pxd b/av/filter/graph.pxd index c9226749c..b3bf352a3 100644 --- a/av/filter/graph.pxd +++ b/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 diff --git a/av/filter/graph.pyi b/av/filter/graph.pyi index 75930d08a..625364f35 100644 --- a/av/filter/graph.pyi +++ b/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: ... diff --git a/av/filter/graph.pyx b/av/filter/graph.pyx index f6376b3a3..53689c432 100644 --- a/av/filter/graph.pyx +++ b/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() diff --git a/tests/test_filters.py b/tests/test_filters.py index 326ef7fb9..3d9e6e9d0 100644 --- a/tests/test_filters.py +++ b/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)