diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aed3b41fb..3b9107b1f 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -32,7 +32,8 @@ v19.0.3 (Unreleased) -------------------- Fixes: -- + +- Support slicing ``Container.streams``, which raised ``TypeError`` despite being typed. v19.0.2 ------- diff --git a/av/audio/stream.py b/av/audio/stream.py index 0efde1a06..9256db4ea 100644 --- a/av/audio/stream.py +++ b/av/audio/stream.py @@ -46,6 +46,19 @@ def encode(self, frame: AudioFrame | None = None): return packets + def encode_lazy(self, frame: AudioFrame | None = None): + """ + Encode an :class:`.AudioFrame` and yield each :class:`.Packet` as it is produced. + + .. seealso:: This is mostly a passthrough to :meth:`.CodecContext.encode_lazy`. + """ + self._assert_has_codec_context(lib.AVERROR_ENCODER_NOT_FOUND) + packet: Packet + for packet in self.codec_context.encode_lazy(frame): + packet._stream = self + packet.ptr.stream_index = self.ptr.index + yield packet + @cython.ccall def decode(self, packet: Packet | None = None): """ diff --git a/av/audio/stream.pyi b/av/audio/stream.pyi index f92fb52ba..88dfbc35b 100644 --- a/av/audio/stream.pyi +++ b/av/audio/stream.pyi @@ -1,3 +1,4 @@ +from collections.abc import Iterator from typing import Literal from av.packet import Packet @@ -19,6 +20,7 @@ class _Layout: class AudioStream(Stream): codec_context: AudioCodecContext def encode(self, frame: AudioFrame | None = None) -> list[Packet]: ... + def encode_lazy(self, frame: AudioFrame | None = None) -> Iterator[Packet]: ... def decode(self, packet: Packet | None = None) -> list[AudioFrame]: ... # From codec context diff --git a/av/container/streams.py b/av/container/streams.py index 3b1dca00b..8e0b5f530 100644 --- a/av/container/streams.py +++ b/av/container/streams.py @@ -75,6 +75,8 @@ def __iter__(self): def __getitem__(self, index): if isinstance(index, int): return self.get(index)[0] + elif isinstance(index, slice): + return cython.cast(list, self._streams)[index] else: return self.get(index) diff --git a/av/video/stream.py b/av/video/stream.py index b4d0508cc..e2efe667b 100644 --- a/av/video/stream.py +++ b/av/video/stream.py @@ -53,6 +53,19 @@ def encode(self, frame: VideoFrame | None = None): packet.ptr.stream_index = self.ptr.index return packets + def encode_lazy(self, frame: VideoFrame | None = None): + """ + Encode an :class:`.VideoFrame` and yield each :class:`.Packet` as it is produced. + + .. seealso:: This is mostly a passthrough to :meth:`.CodecContext.encode_lazy`. + """ + self._assert_has_codec_context(lib.AVERROR_ENCODER_NOT_FOUND) + packet: Packet + for packet in self.codec_context.encode_lazy(frame): + packet._stream = self + packet.ptr.stream_index = self.ptr.index + yield packet + @cython.ccall def decode(self, packet: Packet | None = None): """ diff --git a/tests/test_encode.py b/tests/test_encode.py index f9fedc8b6..877f7f060 100644 --- a/tests/test_encode.py +++ b/tests/test_encode.py @@ -330,6 +330,38 @@ def test_stream_index(self) -> None: assert apacket.stream is astream assert apacket.stream_index == 1 + def test_stream_index_lazy(self) -> None: + with av.open(self.sandboxed("output_lazy.mov"), "w") as output: + astream = output.add_stream("mp2", 48000) + astream.layout = "stereo" + astream.format = "s16" + + vstream = output.add_stream("mpeg4", 24) + vstream.pix_fmt = "yuv420p" + vstream.width = 320 + vstream.height = 240 + + assert astream.index == 0 + assert vstream.index == 1 + + vpackets = list(vstream.encode_lazy(VideoFrame(320, 240, "yuv420p"))) + vpackets += list(vstream.encode_lazy(None)) + assert vpackets + for vpacket in vpackets: + assert vpacket.stream is vstream + assert vpacket.stream_index == 1 + + apackets = [] + for i in range(10): + aframe = AudioFrame("s16", "stereo", samples=astream.frame_size or 1000) + aframe.sample_rate = 48000 + apackets += list(astream.encode_lazy(aframe)) + apackets += list(astream.encode_lazy(None)) + assert apackets + for apacket in apackets: + assert apacket.stream is astream + assert apacket.stream_index == 0 + def test_stream_audio_resample(self) -> None: with av.open(self.sandboxed("output.mov"), "w") as output: vstream = output.add_stream("mpeg4", 24) diff --git a/tests/test_streams.py b/tests/test_streams.py index 187f37241..4a47a70ff 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -124,6 +124,8 @@ def test_selection(self) -> None: assert [video] == container.streams.get(video=0) assert [video] == container.streams.get(video=(0,)) + assert [video, audio] == container.streams[:] + assert [audio] == container.streams[1:] assert video == container.streams.best("video") assert audio == container.streams.best("audio")