Skip to content

Commit

Permalink
feat(client): support passing chunk size for binary responses (#175)
Browse files Browse the repository at this point in the history
  • Loading branch information
stainless-bot committed Nov 15, 2023
1 parent 05ff244 commit e5ebf07
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 6 deletions.
18 changes: 14 additions & 4 deletions src/finch/_base_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
15 changes: 13 additions & 2 deletions src/finch/_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand Down Expand Up @@ -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.
"""
Expand Down

0 comments on commit e5ebf07

Please sign in to comment.