Move transform_output to helpers module

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

Author: mattt

replicate/helpers.py

@@ -1,6 +1,7 @@
 import base64
 import io
 import mimetypes
+from collections.abc import Mapping, Sequence
 from pathlib import Path
 from types import GeneratorType
 from typing import TYPE_CHECKING, Any, Optional
@@ -108,3 +109,20 @@ def base64_encode_file(file: io.IOBase) -> str:
         mimetypes.guess_type(getattr(file, "name", ""))[0] or "application/octet-stream"
     )
     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
+
+    def transform(obj: Any) -> Any:
+        if isinstance(obj, Mapping):
+            return {k: transform(v) for k, v in obj.items()}
+        elif isinstance(obj, Sequence) and not isinstance(obj, str):
+            return [transform(item) for item in obj]
+        elif isinstance(obj, str) and (
+            obj.startswith("https:") or obj.startswith("data:")
+        ):
+            return FileOutput(obj, client)
+        return obj
+
+    return transform(value)

replicate/run.py

@@ -1,4 +1,3 @@
-from collections.abc import Mapping, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -14,10 +13,10 @@
 
 from replicate import identifier
 from replicate.exceptions import ModelError
+from replicate.helpers import transform_output
 from replicate.model import Model
 from replicate.prediction import Prediction
 from replicate.schema import make_schema_backwards_compatible
-from replicate.stream import FileOutput
 from replicate.version import Version, Versions
 
 if TYPE_CHECKING:
@@ -140,19 +139,4 @@ def _make_async_output_iterator(
     return None
 
 
-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()}
-        elif isinstance(obj, Sequence) and not isinstance(obj, str):
-            return [transform(item) for item in obj]
-        elif isinstance(obj, str) and (
-            obj.startswith("https:") or obj.startswith("data:")
-        ):
-            return FileOutput(obj, client)
-        return obj
-
-    return transform(value)
-
-
 __all__: List = []

tests/test_run.py

@@ -1,11 +1,11 @@
 import asyncio
 import sys
+from typing import cast
 
 import httpx
 import pytest
 import respx
 
-from typing import cast
 import replicate
 from replicate.client import Client
 from replicate.exceptions import ModelError, ReplicateError