Move FileOutput to helpers module

Signed-off-by: Mattt Zmuda <[email protected]>

Author: mattt

replicate/helpers.py

@@ -4,7 +4,9 @@
 from collections.abc import Mapping, Sequence
 from pathlib import Path
 from types import GeneratorType
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Optional
+
+import httpx
 
 if TYPE_CHECKING:
     from replicate.client import Client
@@ -111,9 +113,44 @@ def base64_encode_file(file: io.IOBase) -> str:
     return f"data:{mime_type};base64,{encoded_body}"
 
 
-def transform_output(value: Any, client: "Client") -> Any:
-    from replicate.stream import FileOutput  # pylint: disable=import-outside-toplevel
+class FileOutput(httpx.ByteStream, httpx.AsyncByteStream):
+    url: str
+    client: "Client"
+
+    def __init__(self, url: str, client: "Client"):
+        self.url = url
+        self.client = client
+
+    def read(self) -> bytes:
+        with self.client._client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            return response.read()
+
+    def __iter__(self) -> Iterator[bytes]:
+        with self.client._client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            for chunk in response.iter_bytes():
+                yield chunk
+
+    async def aread(self) -> bytes:
+        async with self.client._async_client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            return await response.aread()
 
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        async with self.client._async_client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            async for chunk in response.aiter_bytes():
+                yield chunk
+
+    def __str__(self) -> str:
+        return self.url
+
+    def __repr__(self) -> str:
+        return self.url
+
+
+def transform_output(value: Any, client: "Client") -> Any:
     def transform(obj: Any) -> Any:
         if isinstance(obj, Mapping):
             return {k: transform(v) for k, v in obj.items()}

replicate/stream.py

@@ -1,6 +1,3 @@
-import io
-import base64
-import httpx
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
@@ -12,9 +9,9 @@
     Optional,
     Union,
 )
-from contextlib import asynccontextmanager, contextmanager
-from typing_extensions import Unpack
 
+import httpx
+from typing_extensions import Unpack
 
 from replicate import identifier
 from replicate.exceptions import ReplicateError
@@ -33,43 +30,6 @@
     from replicate.version import Version
 
 
-class FileOutput(httpx.ByteStream, httpx.AsyncByteStream):
-    url: str
-    client: "Client"
-
-    def __init__(self, url: str, client: "Client"):
-        self.url = url
-        self.client = client
-
-    def read(self) -> bytes:
-        with self.client._client.stream("GET", self.url) as response:
-            response.raise_for_status()
-            return response.read()
-
-    def __iter__(self) -> Iterator[bytes]:
-        with self.client._client.stream("GET", self.url) as response:
-            response.raise_for_status()
-            for chunk in response.iter_bytes():
-                yield chunk
-
-    async def aread(self) -> bytes:
-        async with self.client._async_client.stream("GET", self.url) as response:
-            response.raise_for_status()
-            return await response.aread()
-
-    async def __aiter__(self) -> AsyncIterator[bytes]:
-        async with self.client._async_client.stream("GET", self.url) as response:
-            response.raise_for_status()
-            async for chunk in response.aiter_bytes():
-                yield chunk
-
-    def __str__(self) -> str:
-        return self.url
-
-    def __repr__(self) -> str:
-        return self.url
-
-
 class ServerSentEvent(pydantic.BaseModel):  # type: ignore
     """
     A server-sent event.

tests/test_run.py

@@ -9,7 +9,7 @@
 import replicate
 from replicate.client import Client
 from replicate.exceptions import ModelError, ReplicateError
-from replicate.stream import FileOutput
+from replicate.helpers import FileOutput
 
 
 @pytest.mark.vcr("run.yaml")