From 739528dd910358b4d45052076a5814f0c2460d93 Mon Sep 17 00:00:00 2001
From: Stainless Bot <107565488+stainless-bot@users.noreply.github.com>
Date: Wed, 8 Nov 2023 22:27:27 +0000
Subject: [PATCH] feat(client): support passing chunk size for binary responses
 (#747)

---
 src/openai/_base_client.py | 18 ++++++++++++++----
 src/openai/_types.py       | 15 +++++++++++++--
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/openai/_base_client.py b/src/openai/_base_client.py
index e37759cdf8..b2fe242634 100644
--- a/src/openai/_base_client.py
+++ b/src/openai/_base_client.py
@@ -1727,9 +1727,14 @@ def iter_raw(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
         return self.response.iter_raw(chunk_size)
 
     @override
-    def stream_to_file(self, file: str | os.PathLike[str]) -> None:
+    def stream_to_file(
+        self,
+        file: str | os.PathLike[str],
+        *,
+        chunk_size: int | None = None,
+    ) -> None:
         with open(file, mode="wb") as f:
-            for data in self.response.iter_bytes():
+            for data in self.response.iter_bytes(chunk_size):
                 f.write(data)
 
     @override
@@ -1757,10 +1762,15 @@ async def aiter_raw(self, chunk_size: Optional[int] = None) -> AsyncIterator[byt
         return self.response.aiter_raw(chunk_size)
 
     @override
-    async def astream_to_file(self, file: str | os.PathLike[str]) -> None:
+    async def astream_to_file(
+        self,
+        file: str | os.PathLike[str],
+        *,
+        chunk_size: int | None = None,
+    ) -> None:
         path = anyio.Path(file)
         async with await path.open(mode="wb") as f:
-            async for data in self.response.aiter_bytes():
+            async for data in self.response.aiter_bytes(chunk_size):
                 await f.write(data)
 
     @override
diff --git a/src/openai/_types.py b/src/openai/_types.py
index dabd15866f..0d05be9493 100644
--- a/src/openai/_types.py
+++ b/src/openai/_types.py
@@ -123,7 +123,12 @@ def iter_raw(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
         pass
 
     @abstractmethod
-    def stream_to_file(self, file: str | PathLike[str]) -> None:
+    def stream_to_file(
+        self,
+        file: str | PathLike[str],
+        *,
+        chunk_size: int | None = None,
+    ) -> None:
         """
         Stream the output to the given file.
         """
@@ -172,7 +177,13 @@ async def aiter_raw(self, chunk_size: Optional[int] = None) -> AsyncIterator[byt
         """
         pass
 
-    async def astream_to_file(self, file: str | PathLike[str]) -> None:
+    @abstractmethod
+    async def astream_to_file(
+        self,
+        file: str | PathLike[str],
+        *,
+        chunk_size: int | None = None,
+    ) -> None:
         """
         Stream the output to the given file.
         """