From 0fa5e53abd884f814d40486dca270f5dbc684f47 Mon Sep 17 00:00:00 2001
From: WyattBlue <wyattblue@auto-editor.com>
Date: Wed, 21 Oct 2026 18:21:53 -0400
Subject: [PATCH] Add `encode_lazy` to `VideoStream` and `AudioStream`

Previously `stream.encode_lazy()` was forwarded to the codec context via
`__getattr__`, so packets never had their stream or stream_index set and
were muxed into stream 0 regardless of which stream encoded them.
---
 av/audio/stream.py   | 13 +++++++++++++
 av/audio/stream.pyi  |  2 ++
 av/video/stream.py   | 13 +++++++++++++
 tests/test_encode.py | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 60 insertions(+)

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/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)