✨ ⏫ way simpler remotebox

shroominic · shroominic · commit 46fed727cb86 · 2024-07-17T13:31:40.000-07:00
diff --git a/src/codeboxapi/remote.py b/src/codeboxapi/remote.py
@@ -1,84 +1,75 @@
-from json import loads
-from os import PathLike
-from typing import AsyncGenerator, AsyncIterator, BinaryIO, Generator, Iterator, Literal
+from os import PathLike, getenv
+from typing import AsyncGenerator, BinaryIO, Generator, Literal
 from uuid import uuid4
 
-import anyio
 import httpx
 
-from . import utils
 from .codebox import CodeBox, CodeBoxFile, ExecChunk
-from .config import settings
+from .utils import raise_error, resolve_pathlike
 
 
 class RemoteBox(CodeBox):
     """
     Sandboxed Python Interpreter
     """
 
-    def __new__(cls, *args, **kwargs):
-        if kwargs.pop("local", False) or settings.api_key == "local":
-            from .local import LocalBox
-
-            return LocalBox(*args, **kwargs)
-        return super().__new__(cls)
+    def __new__(cls) -> "RemoteBox":
+        # This is a hack to ignore the CodeBox.__new__ factory method.
+        return object.__new__(cls)
 
     def __init__(
         self,
-        factory_id: str | None = None,
-        api_key: str | None = None,
+        session_id: str | None = None,
+        api_key: str | Literal["local", "docker"] = "local",
+        factory_id: str | Literal["default"] = "default",
+        base_url: str = "https://codeboxapi.com/api/v2",
+        _new: bool = False,
     ) -> None:
-        super().__init__()
-        self.session_id = uuid4().hex
+        self.session_id = session_id or uuid4().hex
         self.factory_id = factory_id
-        self.api_key = api_key or settings.api_key
-        self.aclient = httpx.AsyncClient(
-            base_url=f"{settings.base_url}/codebox/{self.session_id}"
+        self.api_key = (
+            api_key
+            or getenv("CODEBOX_API_KEY")
+            or raise_error("CODEBOX_API_KEY is required")
         )
+        self.base_url = f"{base_url}/codebox/{self.session_id}"
+        self.headers = {"Factory-Id": self.factory_id} if self.factory_id else None
+        self.client = httpx.Client(base_url=self.base_url, headers=self.headers)
+        self.aclient = httpx.AsyncClient(base_url=self.base_url, headers=self.headers)
 
     def stream_exec(
         self,
         code: str | PathLike,
-        language: Literal["python", "bash"] = "python",
+        kernel: Literal["ipython", "bash"] = "ipython",
         timeout: float | None = None,
         cwd: str | None = None,
     ) -> Generator[ExecChunk, None, None]:
-        async_gen = self.astream_exec(code, language, timeout, cwd)
-        return (chunk for chunk in anyio.run(utils.collect_async_gen, async_gen))
-
-    def upload(
-        self,
-        file_name: str,
-        content: BinaryIO | bytes | str,
-        timeout: float | None = None,
-    ) -> CodeBoxFile:
-        return anyio.run(self.aupload, file_name, content, timeout)
-
-    def stream_download(
-        self,
-        remote_file_path: str,
-        timeout: float | None = None,
-    ) -> Iterator[bytes]:
-        return anyio.run(
-            utils.collect_async_gen, self.astream_download(remote_file_path, timeout)
-        )
+        code = resolve_pathlike(code)
+        with self.client.stream(
+            method="POST",
+            url="/stream",
+            timeout=timeout,
+            params={"code": code, "kernel": kernel, "cwd": cwd},
+        ) as response:
+            for chunk in response.iter_text():
+                yield ExecChunk.decode(chunk)
 
     async def astream_exec(
         self,
         code: str | PathLike,
-        language: Literal["python", "bash"] = "python",
+        kernel: Literal["ipython", "bash"] = "ipython",
         timeout: float | None = None,
         cwd: str | None = None,
     ) -> AsyncGenerator[ExecChunk, None]:
-        code = utils.resolve_pathlike(code)
+        code = resolve_pathlike(code)
         async with self.aclient.stream(
             method="POST",
             url="/stream",
             timeout=timeout,
-            params={"code": code, "language": language, "cwd": cwd},
+            params={"code": code, "kernel": kernel, "cwd": cwd},
         ) as response:
             async for chunk in response.aiter_text():
-                yield ExecChunk(**loads(chunk))
+                yield ExecChunk.decode(chunk)
 
     async def aupload(
         self,
@@ -96,14 +87,28 @@ async def aupload(
                     timeout=timeout,
                 )
             ).json(),
-            codebox=self,
+            codebox_id=self.session_id,
         )
 
+    def stream_download(
+        self,
+        remote_file_path: str,
+        timeout: float | None = None,
+    ) -> Generator[bytes, None, None]:
+        with self.client.stream(
+            method="GET",
+            url="/download",
+            timeout=timeout,
+            params={"file_name": remote_file_path},
+        ) as response:
+            for chunk in response.iter_bytes():
+                yield chunk
+
     async def astream_download(
         self,
         remote_file_path: str,
         timeout: float | None = None,
-    ) -> AsyncIterator[bytes]:
+    ) -> AsyncGenerator[bytes, None]:
         async with self.aclient.stream(
             method="GET",
             url="/download",
